[leetcode] 2049 统计最高分的节点数目 | dfs二叉树

class Solution {
public:
    int cnt[100000 + 1];
    vector<int> son[100000 + 1];
    int countHighestScoreNodes(vector<int> &parents) {
        const int siz = parents.size();
        if (siz == 2) return 2;
        memset(cnt, 0, sizeof cnt);
        int root;
        for (int i = 0; i < siz; i++) {
            int fa = parents[i];
            if (fa == -1) {
                root = i;
                continue;
            }
            son[fa].push_back(i);
        }
        cnt[root] = dfs(root);
        long long mx = 0, ret = 0;
        for (int i = 0; i < siz; i++) {
            if (root == i) {
                long long tmx = 1;
                for (int x: son[i]) tmx *= cnt[x];
                if (tmx < mx) continue;
                else if (tmx == mx) ret++;
                else {
                    mx = tmx;
                    ret = 1;
                }
            } else {
                long long tmx = 1;
                for (int x: son[i]) tmx *= cnt[x];
                tmx *= siz - cnt[i];
                if (tmx < mx) continue;
                else if (tmx == mx) ret++;
                else {
                    mx = tmx;
                    ret = 1;
                }
            }
            // cout << i << "   " << mx << endl;
        }
        return ret;
    }
    int dfs(int root) {
        int sum = 0;
        for (int i = 0; i < son[root].size(); i++) {
            int u = son[root][i];
            cnt[u] = dfs(u);
            sum += cnt[u];
        }
        return sum + 1;
    }
};