Count Nodes With the Highest Score LeetCode solution in C++

1
class Solution {
2
 public:
3
  // Steps :
4
  // 1 - For each node, you need to find the sizes of the subtrees rooted in
5
  // each of its children.
6

7
  // 2 - How to determine the number of nodes in the rest of the tree?
8
  // Can you subtract the size of the subtree rooted at the node from the total
9
  // number of nodes of the tree?
10

11
  // 3 - Use these values to compute the score of the node. Track the maximum
12
  // score, and how many nodes achieve such score.
13

14
  // calculating size of each subtree by standing at every node  '0' to 'n-1'
15
  int dfs(int src, vector<vector<int>> &g, vector<int> &size) {
16
    int ans = 1; // for curent node
17
    for (auto child : g[src]) {
18
      ans += dfs(child, g, size);
19
    }
20
    return size[src] = ans;
21
  }
22

23
  // This code can also be work for generalized tree not only for Binary tree
24
  int countHighestScoreNodes(vector<int> &parents) {
25
    int n = parents.size();
26
    vector<int> size(n, 0);   // size[i] indicates size of subtree(rooted at i node) + 1
27
    vector<vector<int>> g(n); // storing left and right child of a node
28
    for (int i = 1; i < n; i++) {
29
      g[parents[i]].push_back(i); // 'There is no parent for 0th node'
30
    }
31

32
    dfs(0, g, size); // calculating size of each subtree(rooted at ith node)
33

34
    long long int maxCount = 0; // To avoid overflow  because perform product
35
                                // below you should take "long long int"
36
    long long int maxScore = 0;
37
    for (int i = 0; i < n; i++) // Nodes from '0' to 'n-1'
38
    {
39
      // calculate score of each node after removal their 'edge' or 'node
40
      // itself'.
41
      long long int product = 1;
42
      product = max(product,
43
                    (long long int)(n - size[i])); // calculating leftover nodes
44
                                                   // excluding child nodes
45
      for (auto x : g[i]) {
46
        product = product * size[x];
47
      }
48

49
      if (product > maxScore) {
50
        maxScore = product;
51
        maxCount = 1;
52
      } else if (product == maxScore) {
53
        maxCount++; // store count of nodes which have maximum score equal to
54
                    // "maxScore"
55
      }
56
    }
57

58
    return maxCount;
59
  }
60
};