2049. Count Nodes With the Highest Score

Problem Description

This problem presents a binary tree rooted at node 0 and consisting of n nodes, each uniquely labeled from 0 to n - 1. We are given an integer array parents where each element represents the parent of a corresponding node, such that parents[i] is the parent of node i. The root node’s parent is indicated by parents[0] being -1.

The score of a node is determined by removing the node along with its connected edges from the tree. This operation results in a number of non-empty subtrees. The size of each subtree is simply the number of nodes it contains. The score for the removed node is calculated as the product of the sizes of all resultant subtrees. The goal is to find out how many nodes in the tree have the highest score.

Intuition

To solve this problem, we need a way to compute the score of each node efficiently. Instead of actually removing nodes, which would be inefficient, we can simulate the process with a depth-first search (DFS) algorithm. Here’s the logic behind the solution:

• A tree structure often calls for a recursive strategy. DFS is a good fit here as we need to explore subtrees to calculate scores.
• The solution tracks the size of each subtree as the DFS traversal visits the nodes. The size of the subtree rooted at a node is 1 (for the node itself) plus the sizes of all subtrees of its children.
• If a node is removed, the score for that node is the product of the sizes of the subtrees left behind. However, there is one additional "subtree" to consider — the one formed by the rest of the tree outside of this node's subtree. Its size is the total number of nodes in the tree minus the size of the subtree rooted at the given node.
• For the root node, its score is simply the product of the sizes of its child subtrees because there is no "parent" subtree to consider.
• In the recursive DFS function, we track the maximum score found and count how many times nodes with this score occur.
• For each node visited, we calculate a temporary score. If it's greater than the current maximum score, we update the maximum score and reset our count of nodes with the maximum score to 1. If the score matches the current maximum, we increment our count.
• After the DFS completes, the counter will have the number of nodes with the highest score.

Putting it all together, starting the DFS from the root will give us the needed information to return the final count of nodes with the highest score.

Learn more about Tree, Depth-First Search and Binary Tree patterns.

Solution Approach

The solution employs a Depth-First Search (DFS) algorithm to compute subtree sizes and uses a simple list (Python's List data structure) to model the graph (tree) structure.

1. Transformation into a graph: The given parents array is used to create an adjacency list that represents the tree. For each node (excluding the root), we append it to a list at the index of its parent. This means g[i] will contain a list of all children of the ith node.

2. DFS algorithm: We define a recursive function dfs that takes a single argument cur, which represents the current node being visited. The function initializes a local variable size to 1 (for the current node) and score to 1 (the initial product of subtree sizes).

3. Calculating subtree sizes and scores: For every child c of the current node cur, we recursively call dfs(c), which returns the size of the subtree rooted at c. This value is added to the size of the current subtree and also multiplied into the score for the current node.

4. Handling of the non-subtree part: If the current node is not the root, we need to consider the rest of the tree outside of the current node's subtree. We multiply the current score by n - size to incorporate this.

5. Tracking the maximum score and count: We use global variables max_score and ans to keep track of the maximum score found so far, and the number of nodes with that score, respectively. If the score of the current node is higher than the max_score, we update max_score and set ans to 1. If the score is equal to max_score, we increment ans.

6. Starting the DFS and returning the result: The DFS is started by calling dfs(0) because the root is labeled 0. Once DFS is completed, ans will hold the count of nodes with the highest score, and this is what is returned.

The elegance of the solution comes from its efficient traversal of the tree using DFS to calculate the scores for all nodes without any modifications to the tree's structure, and clever use of global variables to keep track of the maximum score observed, as well as the count of nodes that achieve this score.

Example Walkthrough

To illustrate the solution approach, let’s consider a small tree with n = 5 nodes, and the following parents array representing the tree:

1Node index:  0   1   2   3   4
2Parents:    [-1,  0,  0,  1,  1]

This represents a tree where:

• Node 0 is the root (as parents[0] = -1).
• Nodes 1 and 2 are children of node 0.
• Nodes 3 and 4 are children of node 1.

First, we build the adjacency list based on the parents array to represent the tree graph:

1g[0] -> [1, 2]
2g[1] -> [3, 4]
3g[2] -> []
4g[3] -> []
5g[4] -> []

Now we run the DFS algorithm starting from the root (node 0). We also initialize our global variables max_score and ans to track the maximum score and number of nodes with that maximum score.

1. Start DFS at root node 0. Initialize size = 1, score = 1.

2. For every child of node 0, that is nodes 1 and 2, we call dfs(child).

3. Calling dfs(1):

   • Start with size = 1, score = 1.
   • Visit children of node 1, which are nodes 3 and 4.
   • Call dfs(3), which returns 1 because it has no children, adding this to size of node 1 and setting score = score * size_subtree(3) = 1.
   • Call dfs(4), similar to dfs(3), return 1, adding to size of node 1, and score = score * size_subtree(4) = 1.
   • Final size of subtree at node 1 is 3, and the score for node 1 when removing it would be score * (n - size) = 1 * (5 - 3) = 2. Update max_score to 2 and ans to 1.

4. Back to node 0, now calling dfs(2):

   • Node 2 has no children, so it returns a size of 1.
   • Add 1 to size of node 0 which is now 5 (the total size of the tree).

5. With DFS complete for node 0, node 0 score would be the product of subtree sizes of its children, which are nodes 1 and 2. So score = size_subtree(1) * size_subtree(2) = 3 * 1 = 3. The max_score is less than 3, so we update max_score to 3 and reset ans to 1.

6. DFS has now visited all nodes. Since no other nodes will have a score higher than 3 (the score when the root is removed), we conclude that the highest score is 3, and there is only 1 node (the root, node 0) with this score.

Based on the solution approach, we return ans, which is 1, indicating there is one node (the root) with the highest score in the tree.

Solution Implementation

Time and Space Complexity

Time Complexity

The time complexity of the code is O(n), where n is the number of nodes in the tree. The code uses a Depth-First Search (DFS) traversal to compute the size and score for each node exactly once. For each node, it looks into its children and calculates the score based on the size of the subtrees which are the results of the DFS calls. These operations all happen within the DFS function which is called once per node, leading to a linear time complexity with respect to the number of nodes.

Space Complexity

The space complexity of the code can also be considered O(n) for several reasons:

• The adjacency list g, which represents the tree, will contain a maximum of n-1 edges, as it is a tree.
• The recursive DFS will at most be called recursively for a depth equal to the height of the tree. In the worst case (a skewed tree), the recursive call stack could also be O(n).
• The auxiliary space required by the internal state of the DFS (variables like size, score, and the returned values) will not exceed O(1) per call.

But, typically, the tree height is much less than n for a reasonably balanced tree, so the average case for the space complexity due to recursive calls is less than O(n). However, since the worst-case scenario can still be O(n) (for a skewed tree), we mention it as such.

Learn more about how to find time and space complexity quickly using problem constraints.