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.

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.

Python Solution

1from typing import List
2
3class Solution:
4    def countHighestScoreNodes(self, parents: List[int]) -> int:
5        # Initialize the number of nodes, maximum score, and answer counter.
6        num_nodes = len(parents)
7        max_score = 0
8        answer = 0
9      
10        # Graph representation of the tree, with each node pointing to its children.
11        graph = [[] for _ in range(num_nodes)]
12      
13        # Build the adjacency list for each parent.
14        for i in range(1, num_nodes):
15            graph[parents[i]].append(i)
16      
17        # Depth-First Search to compute subtree sizes and scores.
18        def dfs(current_node: int) -> int:
19            nonlocal max_score, answer
20            subtree_size = 1
21            node_score = 1
22          
23            # Visit each child and calculate the score contribution of its subtree.
24            for child in graph[current_node]:
25                child_subtree_size = dfs(child)
26                subtree_size += child_subtree_size
27                node_score *= child_subtree_size
28          
29            # If not the root, multiply by the size of the "complement" subtree.
30            if current_node > 0:
31                node_score *= num_nodes - subtree_size
32          
33            # Update answer and max_score in case of new max or equal score found.
34            if node_score > max_score:
35                max_score = node_score
36                answer = 1
37            elif node_score == max_score:
38                answer += 1
39          
40            return subtree_size
41      
42        # Start DFS from the root of the tree (node 0).
43        dfs(0)
44      
45        # Return the number of nodes that have the maximum score.
46        return answer
47

Java Solution

1class Solution {
2
3    private int nodeCount; // Total number of nodes in the tree
4    private long maximumScore; // The highest score found so far
5    private int countOfMaxScoreNodes; // The count of nodes having the maximum score
6    private List<List<Integer>> childrenList; // Adjacency list representation of the tree
7
8    public int countHighestScoreNodes(int[] parents) {
9        nodeCount = parents.length;
10        maximumScore = 0;
11        countOfMaxScoreNodes = 0;
12        childrenList = new ArrayList<>();
13      
14        // Initialize lists to hold children nodes for each node
15        for (int i = 0; i < nodeCount; i++) {
16            childrenList.add(new ArrayList<>());
17        }
18      
19        // Build the adjacency list (tree graph) from the parent array
20        for (int i = 1; i < nodeCount; i++) {
21            childrenList.get(parents[i]).add(i);
22        }
23      
24        // Start Depth-First Search from the root node (0)
25        dfs(0);
26      
27        // Return the count of nodes that have the highest score
28        return countOfMaxScoreNodes;
29    }
30  
31    // A method to perform a DFS and calculate the size and score of the current subtree
32    private int dfs(int currentNode) {
33        int subTreeSize = 1; // Current node's size is at least 1 (itself)
34        long score = 1; // Initialize score for the current node
35      
36        // Iterate through the children of the current node
37        for (int child : childrenList.get(currentNode)) {
38            int childSubTreeSize = dfs(child); // Recursively get child subtree size
39          
40            // Accumulate total subtree size and compute the score contribution of this child
41            subTreeSize += childSubTreeSize;
42            score *= childSubTreeSize;
43        }
44      
45        // For non-root nodes, multiply the score with the size of the "rest of the tree"
46        if (currentNode > 0) {
47            score *= (nodeCount - subTreeSize);
48        }
49      
50        // Compare and update the maximum score and the count of nodes
51        if (score > maximumScore) {
52            maximumScore = score;
53            countOfMaxScoreNodes = 1;
54        } else if (score == maximumScore) {
55            countOfMaxScoreNodes++;
56        }
57      
58        // Return the subtree size
59        return subTreeSize;
60    }
61}
62

C++ Solution

1#include <vector>
2#include <unordered_map>
3using namespace std;
4
5class Solution {
6public:
7    // Member variables to keep track of the total count of nodes with the highest score. 
8    // Also to keep record of the current maximum score and the total number of nodes.
9    int countOfMaxScoreNodes;
10    long long maxNodeScore;
11    int numNodes;
12
13    // Function to start the process and return the number of nodes with the highest score
14    int countHighestScoreNodes(vector<int>& parents) {
15        countOfMaxScoreNodes = 0; // Initialize count of nodes with maximum score to zero
16        maxNodeScore = 0; // Initialize maximum score to zero
17        numNodes = parents.size(); // Set total number of nodes
18        unordered_map<int, vector<int>> graph; // Create a graph from the parent array
19      
20        // Building the graph from the parent-child relationships
21        for (int i = 1; i < numNodes; ++i) {
22            graph[parents[i]].push_back(i);
23        }
24      
25        // Start the DFS traversal from the root node, which is node 0
26        dfs(0, graph);
27      
28        // Return the total count of nodes with the maximum score
29        return countOfMaxScoreNodes;
30    }
31
32    // DFS function to calculate the score of each node and the size of the subtree rooted at each node
33    int dfs(int node, unordered_map<int, vector<int>>& graph) {
34        int subtreeSize = 1; // Size of the subtree including the current node
35        long long nodeScore = 1; // Product of the sizes of each subtree
36      
37        // For every child of the current node
38        for (int child : graph[node]) {
39            int childSubtreeSize = dfs(child, graph); // Recursive DFS call
40            subtreeSize += childSubtreeSize; // Update the size of the current subtree
41            nodeScore *= childSubtreeSize; // Update the score by multiplying the sizes of subtrees
42        }
43      
44        // If the node is not the root, multiply node score by the size of the "rest of the tree"
45        if (node > 0) {
46            nodeScore *= (numNodes - subtreeSize);
47        }
48      
49        // Update maxNodeScore and countOfMaxScoreNodes based on the calculated score for this node
50        if (nodeScore > maxNodeScore) { // Found a new maximum score
51            maxNodeScore = nodeScore;
52            countOfMaxScoreNodes = 1; // Reset count because we found a new maximum
53        } else if (nodeScore == maxNodeScore) {
54            ++countOfMaxScoreNodes; // Increment count because we found another node with the maximum score
55        }
56      
57        // Return the total size of the subtree rooted at the current node
58        return subtreeSize;
59    }
60};
61

Typescript Solution

1/**
2 * This function counts the number of nodes in a tree that have the highest score. 
3 * The score of each node is calculated as the product of the sizes of its subtrees, 
4 * and if the node is not the root, then the size of the remaining tree is also considered.
5 * 
6 * @param parents - an array where `parents[i]` is the parent of the `i`th node.
7 * @returns the number of nodes with the highest score.
8 */
9function countHighestScoreNodes(parents: number[]): number {
10    // The number of nodes in the tree
11    const nodeCount = parents.length;
12
13    // An adjacency list to store the tree, with each index representing a node and storing its children.
14    let edges = Array.from({ length: nodeCount }, () => []);
15
16    // Constructing the tree
17    for (let i = 0; i < nodeCount; i++) {
18        const parent = parents[i];
19        if (parent !== -1) {
20            edges[parent].push(i);
21        }
22    }
23
24    // Initialize the count of nodes with the maximum score and the maximum score itself
25    let maxScoreNodeCount = 0;
26    let maxScore = 0;
27
28    /**
29     * A depth-first search function that calculates the size of the subtree rooted at this index
30     * and the score for the current node.
31     * 
32     * @param index - The current node index we are calculating size and score for
33     * @returns The size of the subtree rooted at the given index
34     */
35    function dfs(index: number): number {
36        let subtreeSize = 1;
37        let nodeScore = 1;
38
39        // Traversing children to calculate score and subtree sizes
40        for (const childIndex of edges[index]) {
41            const childSubtreeSize = dfs(childIndex);
42            subtreeSize += childSubtreeSize;
43            nodeScore *= childSubtreeSize;
44        }
45
46        // If the node is not root, the remaining tree size is also part of the score
47        if (index > 0) {
48            nodeScore *= nodeCount - subtreeSize;
49        }
50
51        // If the calculated score for the current node is higher than the recorded max, update it
52        if (nodeScore > maxScore) {
53            maxScore = nodeScore;
54            maxScoreNodeCount = 1;
55        } else if (nodeScore === maxScore) {
56            // If the current node has a score equal to the max, increment the count
57            maxScoreNodeCount++;
58        }
59
60        return subtreeSize;
61    }
62
63    // Start DFS from the root node (index 0)
64    dfs(0);
65
66    return maxScoreNodeCount;
67}
68

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.


Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


TA 👨‍🏫