LeetCode 6240. 树上最大得分和路径

LeetCode 6240. 树上最大得分和路径

DFS

const int N = 1e5 + 10, M = 2 * N;

int h[N], w[N], ne[M], e[M], idx = 0;
int bt[N], p[N];
int n;

class Solution {
public:
    void add(int a, int b){
        e[idx] = b, ne[idx] = h[a], h[a] = idx ++;
    }
    int mostProfitablePath(vector<vector<int>>& edges, int bob, vector<int>& amount) {
        int n = amount.size();
        memset(h, -1, sizeof h);
        idx = 0;
        for(int i = 0; i < n; i ++) w[i] = amount[i];

        for(auto& e : edges){
            int a = e[0], b = e[1];
            add(a, b), add(b, a);
        }
        auto dfs1 = [&](auto self, int u, int fa) -> void {

            for(int i = h[u]; ~i; i = ne[i]){
                int j = e[i];
                if(j == fa) continue;
                p[j] = u;
                self(self, j, u);
            }
        };

        dfs1(dfs1, 0, -1);

        memset(bt, -1, sizeof bt);
        int t = 0;
        while(true){
            bt[bob] = t;
            t ++;
            if(!bob) break;
            bob = p[bob];
        }

        auto dfs2 = [&](auto self, int u, int fa, int t) -> int{
            int val = 0;
            if(bt[u] == -1 || t < bt[u]) val = w[u];
            else if(bt[u] == t) val = w[u] / 2;

            int mx = -2e9;
            for(int i = h[u]; ~i; i = ne[i]){
                int j = e[i];
                if(j == fa) continue;
                mx = max(mx, self(self, j, u, t + 1));
            }
            if(mx == -2e9) mx = 0;

            return val + mx;
        };

        return dfs2(dfs2, 0, -1, 0);

    }
};