3772. Maximum Subgraph Score in a Tree

You are given an undirected tree with n nodes, numbered from 0 to n - 1. It is represented by a 2D integer array edges​​​​​​​ of length n - 1, where edges[i] = [a_i, b_i] indicates that there is an edge between nodes a_i and b_i in the tree.

You are also given an integer array good of length n, where good[i] is 1 if the i^th node is good, and 0 if it is bad.

Define the score of a subgraph as the number of good nodes minus the number of bad nodes in that subgraph.

For each node i, find the maximum possible score among all connected subgraphs that contain node i.

Return an array of n integers where the i^th element is the maximum score for node i.

A subgraph is a graph whose vertices and edges are subsets of the original graph.

A connected subgraph is a subgraph in which every pair of its vertices is reachable from one another using only its edges.

class Solution {
    long long dfs(int u, int par, vector<vector<int>>& adj, vector<int>& good, vector<int>& res, int tot) {
        res[u] = good[u] ? 1 : -1;
        for(auto& v : adj[u]) {
            if(v == par) continue;
            auto sub = dfs(v,u,adj,good,res,tot);
            if(sub > 0) res[u] += sub;
        }
        return res[u];
    }
    void dfs2(int u, int par, vector<vector<int>>& adj,  vector<int>& good, vector<int>& res, int parSum) {
        res[u] += parSum;
        parSum += good[u] ? 1 : -1;
        for(auto& v : adj[u]) {
            if(v == par) continue;
            if(res[v] > 0) parSum += res[v];
        }
        for(auto& v : adj[u]) {
            if(v == par) continue;
            int best = parSum;
            if(res[v] > 0) best -= res[v];
            dfs2(v,u,adj,good,res,max(0,best));
        }
    }
public:
    vector<int> maxSubgraphScore(int n, vector<vector<int>>& edges, vector<int>& good) {
        vector<int> res(n);
        vector<vector<int>> adj(n);
        for(auto& e : edges) {
            int u = e[0], v = e[1];
            adj[u].push_back(v);
            adj[v].push_back(u);
        }
        int tot = 2 * accumulate(begin(good), end(good), 0) - n;
        dfs(0,-1,adj,good,res,tot);
        dfs2(0,-1,adj,good,res,0);
        return res;
    }
};