2581. Count Number of Possible Root Nodes

HardTreeDepth-First SearchHash TableDynamic Programming

Problem Description

Alice has an undirected tree consisting of n nodes, labeled from 0 to n - 1. This tree is represented by a 2D integer array named edges, where each element edges[i] consists of two nodes a[i] and b[i] that are connected by an edge. Bob is tasked with finding the root of the tree but only knows the edges.

To find the root, Bob can make several guesses. Each guess is an assumption about who the parent of a particular node is. These guesses are represented by a 2D integer array guesses, with each guesses[j] containing two elements: u[j] and v[j], where Bob guesses that u[j] is the parent of v[j]. Alice, instead of confirming each guess, only tells Bob that at least k of his guesses are true.

The goal is to find the number of possible nodes that can be the root of the tree based on Bob’s guesses and Alice's confirmation that at least k of these guesses are true. If no tree can satisfy the condition given the guesses and value of k, the answer should be 0.

Flowchart Walkthrough

To approach the problem using the Flowchart, let's determine the suitable algorithm step-by-step:

Is it a graph?

Yes: The problem involves evaluating relationships between nodes, consistent with a graph structure.

Is it a tree?

Yes: The essence of the problem is to identify possible root nodes within a setup that can be interpreted as a tree (though unformed, the potential for hierarchical relationships suggests a tree structure).

Based on reaching the tree node on the flowchart and affirming it's a tree, the next directed step is:

DFS:

Depth-First Search (DFS) is appropriate as the problem likely involves exploring node connections deeply to identify and count potential roots, making DFS the suited algorithm from this point in the tree structure of the flowchart.

Conclusion: The flowchart indicates the use of DFS due to the tree-like structure of the problem scenario.

Intuition

To solve this problem, we need to consider each node as a potential root and verify if at least k guesses are true when that node is the root. The core idea is to use Depth-First Search (DFS) to explore each node starting from an arbitrary node (in this case, node 0) and count the number of correct guesses associated with the path from the current node to all its descendants. This number is stored in a variable cnt.

We start the process by considering node 0 as the root and use DFS (dfs1) to calculate the count cnt of correct guesses reachable from node 0. After we have the initial count of correct guesses for the subtree with the root 0, we need to examine whether other nodes can also serve as the root while maintaining at least k correct guesses (dfs2).

To achieve this, as we traverse each node, we adjust the count cnt of correct guesses by:

Subtracting 1 from cnt if the current edge is in Bob's guesses but is not in the correct direction for the new potential root.
Adding 1 to cnt if the edge connecting the current node to its parent is a correct guess in the reversed direction for the current subtree's root.

After the adjustment, if cnt is greater than or equal to k, it means the current node could be a root, so we count that as a possibility. We perform this process for every node by changing the root of the subtree during the DFS traversal, allowing us to identify all possible root nodes that satisfy Alice's condition of at least k correct guesses.

The provided solution uses a hash map gs to represent Bob's guesses for quick lookup during the count adjustment, and the final answer ans is incremented each time a valid root is encountered. DFS ensures that all nodes are visited, and potential roots are evaluated in a single pass through the tree.

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

Solution Approach

The solution relies on a tree Depth-First Search (DFS) algorithm to walk through the tree and count the number of guesses that are correct. There are two DFS methods implemented here, dfs1 and dfs2, both working together to figure out all possible nodes that can be the root.

Initially, we convert the edges list to an adjacency list g which maps each node to its connected nodes (children and parents), enabling easy traversal. This is a common pattern when working with graph-based structures such as trees.
"Tree DP" mentioned in the reference solution approach refers to Tree Dynamic Programming, which is used to cleverly cache results that can be reused to avoid redundant calculations.
A hash map gs (Counter object in Python) is created from the guesses to quickly check if a guess (a directed edge) exists and to maintain counts efficiently. The key in gs is a tuple (u, v), representing a guess where u is assumed to be the parent of v.
dfs1 is used initially to set the baseline count cnt of correct guesses assuming the root is node 0. As dfs1 traverses the tree, it increments cnt if it encounters an edge in gs.
After establishing the baseline cnt, dfs2 traverses the tree again. As it moves from a node i to an adjacent node j, it adjusts cnt to reflect the number of correct guesses if node j were the root. Essentially, it changes the perspective of the root for each subtree, recalculating the guesses accordingly. The rule is:
- cnt -= gs[(i, j)]: If the "parent-to-child" guess (i, j) was previously counted as correct for node i, it must be subtracted when considering node j as the root since (i, j) would no longer represent a "parent-to-child" relation.
- cnt += gs[(j, i)]: Conversely, if the "child-to-parent" guess (j, i) exists, it is now a correct "parent-to-child" relation in the context of node j being the root, hence cnt is incremented.
Following the adjustment, if cnt is greater than or equal to k, it indicates that node j can be a possible root. We use the ans variable to keep track of the number of nodes that meet this condition.
This process is repeated for each node by recursively calling dfs2 as the traversal moves through the tree.

Using DFS and these adjustments, we can cover all nodes and their subtrees, ensuring that the potential roots are counted without redundant recalculations. The hash map gs ensures quick adjustments to the cnt as the DFS progresses, avoiding the need to recount guesses from scratch. The final value of ans will be the total number of possible roots found, which is returned as the solution.

Ready to land your dream job?

Unlock your dream job with a 2-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's take a small example to illustrate the solution approach. Suppose we have a tree with n = 4 nodes and the tree structure is defined by the following edges:

edges = [[0, 1], [0, 2], [2, 3]]

This means node 0 is connected to nodes 1 and 2, and node 2 is connected to node 3. Now, let's imagine Bob makes the following two guesses:

guesses = [[0, 2], [2, 3]]

Here, Bob guesses that node 0 is the parent of node 2, and node 2 is the parent of node 3. Alice tells him that at least k=1 of his guesses are true.

According to the solution approach, we would first build an adjacency list from the edges:

g = {
  0: [1, 2],
  1: [0],
  2: [0, 3],
  3: [2]
}

Then, we create a hash map gs from Bob's guesses:

gs = {
  (0, 2): 1,
  (2, 3): 1
}

We begin DFS with dfs1 starting from node 0, which we arbitrarily assume to be the root:

From 0, we explore its children 1 and 2. The guess (0, 2) exists, so we increment cnt to 1.
There are no guesses involving node 1, so visiting that node does not change cnt.
From node 2, we visit 3, and since the guess (2, 3) is present, we increment cnt to 2.

Now we know that there are 2 correct guesses in the subtree with 0 as the root, which is more than k=1. Therefore, node 0 is a potential root.

Next, we execute dfs2 to check other nodes:

While exploring from node 2, we consider it as a potential root. We adjust cnt:
- We subtract 1 for the edge (0, 2) because if 2 were the root, then (0, 2) would be wrong, so cnt becomes 1.
- We add 1 for the edge (2, 3) which was already counted before. Therefore, cnt stays at 1.

Since cnt is still 1, node 2 could also be a root.

For nodes 1 and 3, there are no incoming guesses, so current cnt will not satisfy k=1 when considering them as roots.

Therefore, the possible node roots based on the minimum k=1 correct guess from Bob would be [0, 2]. The final answer ans would be 2 as there are two possible nodes that could be roots.

Solution Implementation

1from collections import defaultdict, Counter
2from typing import List
3
4class Solution:
5    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:
6        # Helper function for the first depth-first search to calculate the initial count of guesses
7        def dfs_count_initial(node, parent):
8            nonlocal count_guesses
9            for neighbor in graph[node]:
10                # We do not want to go back to the parent node
11                if neighbor != parent:
12                    # Increment count by the number of guesses that this edge has
13                    count_guesses += guesses_count[(node, neighbor)]
14                    dfs_count_initial(neighbor, node)
15      
16        # Helper function for the second depth-first search to calculate the results while moving the root
17        def dfs_collect_results(node, parent):
18            nonlocal result, count_guesses
19            # Increment result if the number of guesses is greater than or equal to k
20            result += count_guesses >= k
21            for neighbor in graph[node]:
22                if neighbor != parent:
23                    # Update count when moving the root up towards this neighbor
24                    count_guesses -= guesses_count[(node, neighbor)]
25                    count_guesses += guesses_count[(neighbor, node)]
26                  
27                    # Recurse into the neighbor
28                    dfs_collect_results(neighbor, node)
29                  
30                    # Backtrack: revert the count changes to search other branches
31                    count_guesses -= guesses_count[(neighbor, node)]
32                    count_guesses += guesses_count[(node, neighbor)]
33      
34        # Graph representation using an adjacency list
35        graph = defaultdict(list)
36        for a, b in edges:
37            graph[a].append(b)
38            graph[b].append(a)
39      
40        # Counter for all the guesses made, represented as a tuple of (node1, node2)
41        guesses_count = Counter(tuple(guess) for guess in guesses)
42      
43        # Initialize the number of guesses count
44        count_guesses = 0
45        dfs_count_initial(0, -1)  # Start DFS from the node labeled 0
46      
47        # Initialize the result which represents the number of trees meeting the condition
48        result = 0
49        dfs_collect_results(0, -1)  # Start DFS again from the node labeled 0 to find results
50      
51        return result  # Return the final result
52

1class Solution {
2    private List<Integer>[] adjacencyList; // Adjacency list to represent the graph
3    private Map<Long, Integer> guessMap = new HashMap<>(); // Map to store the number of guesses between two nodes
4    private int answer; // Variable to store the final answer/result
5    private int threshold; // Threshold k for comparison against the guess count
6    private int guessCount; // Current count of guesses while traversing through the nodes
7    private int nodeCount; // Total number of nodes in the graph
8
9    // Method to find the number of nodes that meet the guess threshold requirement
10    public int rootCount(int[][] edges, int[][] guesses, int k) {
11        this.threshold = k;
12        nodeCount = edges.length + 1;
13        adjacencyList = new List[nodeCount];
14        Arrays.setAll(adjacencyList, e -> new ArrayList<>());
15      
16        // Build the graph as an adjacency list
17        for (var edge : edges) {
18            int from = edge[0], to = edge[1];
19            adjacencyList[from].add(to);
20            adjacencyList[to].add(from);
21        }
22      
23        // Map the guesses into the guessMap with keys created by function 'f' and values being the guess counts
24        for (var guess : guesses) {
25            int from = guess[0], to = guess[1];
26            guessMap.merge(f(from, to), 1, Integer::sum);
27        }
28      
29        // DFS traversal to count the guesses from root to all other nodes
30        dfs1(0, -1);
31      
32        // DFS traversal to count the answers based on guess comparisons to the threshold
33        dfs2(0, -1);
34        return answer;
35    }
36
37    // First depth-first search to count the total guesses from the root
38    private void dfs1(int node, int parent) {
39        for (int child : adjacencyList[node]) {
40            if (child != parent) {
41                guessCount += guessMap.getOrDefault(f(node, child), 0);
42                dfs1(child, node);
43            }
44        }
45    }
46
47    // Second depth-first search to determine the number of nodes that satisfy the threshold condition
48    private void dfs2(int node, int parent) {
49        answer += guessCount >= threshold ? 1 : 0;
50        for (int child : adjacencyList[node]) {
51            if (child != parent) {
52                int guessFromParent = guessMap.getOrDefault(f(node, child), 0);
53                int guessToParent = guessMap.getOrDefault(f(child, node), 0);
54                guessCount -= guessFromParent;
55                guessCount += guessToParent;
56                dfs2(child, node);
57                guessCount -= guessToParent;
58                guessCount += guessFromParent;
59            }
60        }
61    }
62
63    // Function to create a unique key for storing guesses between two nodes in the guessMap
64    private long f(int from, int to) {
65        return ((long) from) * nodeCount + to;
66    }
67}
68

1#include <vector>
2#include <functional>
3#include <unordered_map>
4using namespace std;
5
6class Solution {
7public:
8    int rootCount(vector<vector<int>>& edges, vector<vector<int>>& guesses, int k) {
9        // Calculate the number of nodes
10        int n = edges.size() + 1;
11
12        // Adjacency list to represent the tree
13        vector<vector<int>> graph(n);
14
15        // Map to store the guesses using a pair of nodes as a key
16        unordered_map<long long, int> guessMap; 
17
18        // Helper function to encode two integers into a single long long key
19        auto encode = [&](int i, int j) {
20            return 1LL * i * n + j;
21        };
22
23        // Constructing the graph
24        for (auto& edge : edges) {
25            int a = edge[0], b = edge[1];
26            graph[a].push_back(b);
27            graph[b].push_back(a);
28        }
29
30        // Storing the guesses into the map
31        for (auto& guess : guesses) {
32            int a = guess[0], b = guess[1];
33            guessMap[encode(a, b)]++;
34        }
35
36        int answer = 0; // To store the result
37        int count = 0;  // To store the current count of guessed roots
38
39        // First DFS to count guesses on the path to nodes
40        function<void(int, int)> dfsCount = [&](int node, int parent) {
41            for (int& neighbor : graph[node]) {
42                if (neighbor != parent) {
43                    count += guessMap[encode(node, neighbor)];
44                    dfsCount(neighbor, node);
45                }
46            }
47        };
48
49        // Second DFS to compute the answer while keeping track of count
50        function<void(int, int)> dfsAnswer = [&](int node, int parent) {
51            // If current count is equal or greater than k, increment answer
52            answer += count >= k;
53
54            // Traverse all adjacent nodes
55            for (int& neighbor : graph[node]) {
56                if (neighbor != parent) {
57                    int toNeighborGuesses = guessMap[encode(node, neighbor)];
58                    int fromNeighborGuesses = guessMap[encode(neighbor, node)];
59
60                    // Update the count while moving to the neighbor
61                    count -= toNeighborGuesses;
62                    count += fromNeighborGuesses;
63
64                    // Recurse into neighbor
65                    dfsAnswer(neighbor, node);
66
67                    // Backtrack: restore the count when back from recursion
68                    count -= fromNeighborGuesses;
69                    count += toNeighborGuesses;
70                }
71            }
72        };
73
74        // Launch DFS from the root (node 0) considering it has no parent (-1)
75        dfsCount(0, -1); // First DFS to initialize counts
76        dfsAnswer(0, -1); // Second DFS to compute answer
77
78        // Return the total number of valid roots
79        return answer;
80    }
81};
82

1type Graph = number[][];
2type GuessMap = Map<string, number>;
3
4// Function to encode two integers into a single string key
5const encode = (i: number, j: number, n: number): string => {
6  return `${i * n + j}`;
7};
8
9// Function to calculate the root count given edges, guesses, and k
10const rootCount = (edges: number[][], guesses: number[][], k: number): number => {
11  // Calculate the number of nodes
12  const n: number = edges.length + 1;
13
14  // Adjacency list to represent the treec
15  const graph: Graph = Array.from({ length: n }, () => []);
16
17  // Map to store the guesses with a pair of nodes as a key
18  const guessMap: GuessMap = new Map();
19
20  // Constructing the graph
21  edges.forEach(edge => {
22    const [a, b] = edge;
23    graph[a].push(b);
24    graph[b].push(a);
25  });
26
27  // Storing the guesses in the map
28  guesses.forEach(guess => {
29    const [a, b] = guess;
30    const key: string = encode(a, b, n);
31    const currentCount = guessMap.get(key) || 0;
32    guessMap.set(key, currentCount + 1);
33  });
34
35  let answer: number = 0; // To store the result
36  let count: number = 0;  // To store the current count of guessed roots
37
38  // DFS function to count guesses on the path to nodes
39  const dfsCount = (node: number, parent: number): void => {
40    graph[node].forEach(neighbor => {
41      if (neighbor !== parent) {
42        count += guessMap.get(encode(node, neighbor, n)) || 0;
43        dfsCount(neighbor, node);
44      }
45    });
46  };
47
48  // DFS function to compute the answer while keeping track of count
49  const dfsAnswer = (node: number, parent: number): void => {
50    // If current count is equal or greater than k, increment answer
51    answer += (count >= k) ? 1 : 0;
52
53    graph[node].forEach(neighbor => {
54      if (neighbor !== parent) {
55        const toNeighborGuesses: number = guessMap.get(encode(node, neighbor, n)) || 0;
56        const fromNeighborGuesses: number = guessMap.get(encode(neighbor, node, n)) || 0;
57
58        // Update the count while moving to the neighbor
59        count -= toNeighborGuesses;
60        count += fromNeighborGuesses;
61
62        // Recurse into neighbor
63        dfsAnswer(neighbor, node);
64
65        // Backtrack: restore the count when returning from recursion
66        count -= fromNeighborGuesses;
67        count += toNeighborGuesses;
68      }
69    });
70  };
71
72  // Launch DFS from the root (node 0) considering it has no parent
73  dfsCount(0, -1); // First DFS to initialize counts
74  dfsAnswer(0, -1); // Second DFS to compute the answer
75
76  // Return the total number of valid roots
77  return answer;
78};
79
80export { rootCount }; // Export the function for use in other modules
81

Time and Space Complexity

The time complexity of the algorithm primarily consists of two depth-first search (DFS) operations, dfs1 and dfs2. The dfs1 function is called once for each vertex to calculate the initial count of the guesses in cnt. The dfs2 function is called recursively to traverse all vertices in the graph, incrementing or decrementing cnt as necessary and updating ans.

Both dfs1 and dfs2 visit each vertex once and each edge twice (once from each vertex it connects). Since there are n nodes and each node is connected by one edge, in other words there are n - 1 edges (for a tree). Therefore, each DFS will take O(n) time, totaling O(2n) which is simplified to O(n) for the whole process of DFS.
The construction of graph g and guesses gs will be O(n + m), where n is the length of edges and m is the length of guesses. Each edge and guess is processed only once.

As a result, the overall time complexity is O(n + m), since we have to consider time for both DFS operations and the initial construction of graph structures.

For the space complexity, the algorithm uses additional data structures that store graph information and the counts of guesses:

Graph g stores adjacency lists, which in total will have 2 * (n - 1) entries (each edge is stored twice). This equates to O(n).
Guesses gs are stored as a Counter (a type of dictionary in Python), which in the worst case, will have m unique entries corresponding to each unique guess. This equates to O(m).

Therefore, together with the recursion call stack (which in the worst case, can go as deep as n for a skewed tree), the overall space complexity would be O(n + m), accounting for storing the graph g, the Counter gs, and the depth of the recursion stack.

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

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

Question 1 out of 10

Which type of traversal does breadth first search do?

Pre-order traversal

Post-order traversal

In-order traversal

Level-order traversal

I don’t know