Count Number of Possible Root Nodes

Hard

#2361Time: O(N^2 + M), where N is the number of nodes and M is the number of guesses. Building the adjacency list takes O(N). Storing guesses in a set takes O(M). The main loop runs N times. Inside the loop, a full tree traversal (BFS/DFS) takes O(N) time. So the loop takes O(N*N). The total complexity is O(N^2 + M). This is too slow for the given constraints.Space: O(N + M). O(N) for the adjacency list and the visited array. O(M) for the `guessSet`. The queue for BFS would take O(N) in the worst case.

Patterns

Dynamic Programming

Algorithms

Depth-First Search

Data structures

Array Hash Table Tree

Prompt

Alice has an undirected tree with n nodes labeled from 0 to n - 1. The tree is represented as a 2D integer array edges of length n - 1 where edges[i] = [a_i, b_i] indicates that there is an edge between nodes a_i and b_i in the tree.

Alice wants Bob to find the root of the tree. She allows Bob to make several guesses about her tree. In one guess, he does the following:

Chooses two distinct integers u and v such that there exists an edge [u, v] in the tree.
He tells Alice that u is the parent of v in the tree.

Bob's guesses are represented by a 2D integer array guesses where guesses[j] = [u_j, v_j] indicates Bob guessed u_j to be the parent of v_j.

Alice being lazy, does not reply to each of Bob's guesses, but just says that at least k of his guesses are true.

Given the 2D integer arrays edges, guesses and the integer k, return the number of possible nodes that can be the root of Alice's tree. If there is no such tree, return 0.

Example 1:

Input: edges = [[0,1],[1,2],[1,3],[4,2]], guesses = [[1,3],[0,1],[1,0],[2,4]], k = 3
Output: 3
Explanation: 
Root = 0, correct guesses = [1,3], [0,1], [2,4]
Root = 1, correct guesses = [1,3], [1,0], [2,4]
Root = 2, correct guesses = [1,3], [1,0], [2,4]
Root = 3, correct guesses = [1,0], [2,4]
Root = 4, correct guesses = [1,3], [1,0]
Considering 0, 1, or 2 as root node leads to 3 correct guesses.

Example 2:

Input: edges = [[0,1],[1,2],[2,3],[3,4]], guesses = [[1,0],[3,4],[2,1],[3,2]], k = 1
Output: 5
Explanation: 
Root = 0, correct guesses = [3,4]
Root = 1, correct guesses = [1,0], [3,4]
Root = 2, correct guesses = [1,0], [2,1], [3,4]
Root = 3, correct guesses = [1,0], [2,1], [3,2], [3,4]
Root = 4, correct guesses = [1,0], [2,1], [3,2]
Considering any node as root will give at least 1 correct guess.

Constraints:

edges.length == n - 1
2 <= n <= 10⁵
1 <= guesses.length <= 10⁵
0 <= a_i, b_i, u_j, v_j <= n - 1
a_i != b_i
u_j != v_j
edges represents a valid tree.
guesses[j] is an edge of the tree.
guesses is unique.
0 <= k <= guesses.length

Approaches

2 approaches with complexity analysis and trade-offs.

The most direct way to solve this problem is to test every single node as a potential root. For each node, we can determine the structure of the tree if it were the root, count how many guesses are correct for that structure, and if the count meets the threshold k, we increment our result.

Algorithm

Build an adjacency list representation of the undirected tree from the edges array.\n- Store the guesses in a HashSet for quick lookups. Encode each pair (u, v) into a single long.\n- Initialize a counter for possible roots, result = 0.\n- Loop through each node i from 0 to n-1:\n - Consider i as the current root.\n - Initialize current_correct_guesses = 0.\n - Perform a traversal (BFS or DFS) starting from i. Keep track of visited nodes.\n - During the traversal, for each edge from a parent u to a child v, check if the guess (u, v) is in the guessSet. If yes, increment current_correct_guesses.\n - After the traversal is complete, if current_correct_guesses >= k, increment result.\n- Return result.

Walkthrough

We can iterate through each node from 0 to n-1. Let's say we are currently checking if node r can be a valid root.\n\nTo find out the parent-child relationships for a tree rooted at r, we can perform a tree traversal like Breadth-First Search (BFS) or Depth-First Search (DFS), starting from r. During the traversal, whenever we move from a node u to an unvisited neighbor v, we establish that u is the parent of v.\n\nTo efficiently check the guesses, we can first store all the guesses in a HashSet of pairs. This allows for O(1) lookups. A pair (u, v) can be encoded into a single long value to be stored in the set.\n\nAs we traverse the tree rooted at r and establish a parent-child relationship u -> v, we check if the pair (u, v) exists in our guessSet. If it does, we increment a counter for correct guesses for the root r.\n\nAfter the traversal for root r is complete, we check if the total count of correct guesses is at least k. If it is, we count r as a possible root.\n\nWe repeat this process for all n nodes and return the total count of possible roots.

1import java.util.*;2 3class Solution {4    public int rootCount(int[][] edges, int[][] guesses, int k) {5        int n = edges.length + 1;6        List<List<Integer>> adj = new ArrayList<>();7        for (int i = 0; i < n; i++) {8            adj.add(new ArrayList<>());9        }10        for (int[] edge : edges) {11            adj.get(edge[0]).add(edge[1]);12            adj.get(edge[1]).add(edge[0]);13        }14 15        Set<Long> guessSet = new HashSet<>();16        long factor = n;17        for (int[] guess : guesses) {18            guessSet.add((long) guess[0] * factor + guess[1]);19        }20 21        int possibleRoots = 0;22        for (int i = 0; i < n; i++) {23            // Consider node i as the root24            int correctGuesses = 0;25            26            Queue<Integer> q = new LinkedList<>();27            boolean[] visited = new boolean[n];28            29            q.add(i);30            visited[i] = true;31            32            while (!q.isEmpty()) {33                int u = q.poll();34                for (int v : adj.get(u)) {35                    if (!visited[v]) {36                        visited[v] = true;37                        q.add(v);38                        // u is parent of v39                        if (guessSet.contains((long) u * factor + v)) {40                            correctGuesses++;41                        }42                    }43                }44            }45            46            if (correctGuesses >= k) {47                possibleRoots++;48            }49        }50        51        return possibleRoots;52    }53}

Complexity

Time

O(N^2 + M), where N is the number of nodes and M is the number of guesses. Building the adjacency list takes O(N). Storing guesses in a set takes O(M). The main loop runs N times. Inside the loop, a full tree traversal (BFS/DFS) takes O(N) time. So the loop takes O(N*N). The total complexity is O(N^2 + M). This is too slow for the given constraints.

Space

O(N + M). O(N) for the adjacency list and the visited array. O(M) for the `guessSet`. The queue for BFS would take O(N) in the worst case.

Trade-offs

Pros

Simple to understand and implement.

Cons

Inefficient and will time out on larger test cases due to its quadratic time complexity.

The brute-force approach is inefficient because it recomputes the number of correct guesses from scratch for each potential root. A much more efficient method is the 'rerooting' technique, a form of dynamic programming on trees. We first calculate the answer for an arbitrary root (e.g., node 0), and then in a second pass, we efficiently calculate the answer for every other node by observing how the answer changes when we shift the root to an adjacent node.

Algorithm

Build an adjacency list for the tree.\n- Store guesses in a HashSet for O(1) lookups.\n- DFS Pass 1:\n - Start a DFS from an arbitrary root, say node 0.\n - For each edge u -> v (where u is parent, v is child), check if (u, v) is in the guessSet.\n - Calculate the total number of correct guesses for root 0, let's call it count_root_0.\n- DFS Pass 2:\n - Create an array answer of size n. Set answer[0] = count_root_0.\n - Start a second DFS from root 0.\n - For a traversal from parent u to child v, calculate answer[v] using the formula: answer[v] = answer[u] - (is_guess(u, v) ? 1 : 0) + (is_guess(v, u) ? 1 : 0).\n - Recursively call the DFS for v with its calculated count.\n- Final Count:\n - Iterate through the answer array and count how many values are >= k.

Walkthrough

This approach involves two Depth-First Searches (DFS).\n\nFirst Pass (Post-order traversal style): We arbitrarily choose node 0 as the root. We perform a DFS from node 0 to calculate count0, the number of correct guesses if the tree is rooted at 0. During this DFS, when we traverse from a parent u to a child v, we check if the guess (u, v) exists. We sum up these correct guesses over the entire tree.\n\nSecond Pass (Pre-order traversal style): Now, we perform a second DFS, also starting from node 0. This time, we will calculate the number of correct guesses for every node. Let correct_guesses[i] be the number of correct guesses when i is the root. We already know correct_guesses[0] = count0.\n\nWhen we move the root from a parent u to its child v, the orientation of only one edge changes: the edge between u and v. It was u -> v, and now it becomes v -> u. All other parent-child relationships in the tree remain the same relative to their side of this edge.\n\nTherefore, the number of correct guesses for root v can be calculated from the number of correct guesses for root u with a simple adjustment:\ncorrect_guesses[v] = correct_guesses[u] - (is_guess(u, v) ? 1 : 0) + (is_guess(v, u) ? 1 : 0)\nWe use this formula to propagate the counts down the tree in our second DFS.\n\nAfter the second DFS completes, we will have the correct guess count for every possible root stored in an array. We can then iterate through this array and count how many nodes have a score of at least k.

1import java.util.*;2 3class Solution {4    List<List<Integer>> adj;5    Set<Long> guessSet;6    int n;7    int k;8    int[] correctGuessesCount;9    int countForRoot0 = 0;10    long factor;11 12    public int rootCount(int[][] edges, int[][] guesses, int k) {13        this.n = edges.length + 1;14        this.k = k;15        this.adj = new ArrayList<>();16        for (int i = 0; i < n; i++) {17            adj.add(new ArrayList<>());18        }19        for (int[] edge : edges) {20            adj.get(edge[0]).add(edge[1]);21            adj.get(edge[1]).add(edge[0]);22        }23 24        this.guessSet = new HashSet<>();25        this.factor = n;26        for (int[] guess : guesses) {27            guessSet.add((long) guess[0] * factor + guess[1]);28        }29 30        // Step 1: DFS from root 0 to calculate initial correct guess count31        dfs1(0, -1);32 33        // Step 2: DFS from root 0 again to calculate counts for all other nodes34        this.correctGuessesCount = new int[n];35        dfs2(0, -1, countForRoot0);36 37        // Step 3: Count the number of valid roots38        int result = 0;39        for (int count : correctGuessesCount) {40            if (count >= k) {41                result++;42            }43        }44        return result;45    }46 47    // DFS to calculate the number of correct guesses for the tree rooted at 048    private void dfs1(int u, int p) {49        for (int v : adj.get(u)) {50            if (v == p) continue;51            // Edge is u -> v52            if (guessSet.contains(u * factor + v)) {53                countForRoot0++;54            }55            dfs1(v, u);56        }57    }58 59    // DFS to calculate correct guesses for all nodes using rerooting60    private void dfs2(int u, int p, int currentCorrectGuesses) {61        correctGuessesCount[u] = currentCorrectGuesses;62 63        for (int v : adj.get(u)) {64            if (v == p) continue;65            66            int nextCorrectGuesses = currentCorrectGuesses;67            // The edge u -> v is flipped to v -> u68            // If u -> v was a correct guess, we lose a point69            if (guessSet.contains(u * factor + v)) {70                nextCorrectGuesses--;71            }72            // If v -> u is a guess, we gain a point73            if (guessSet.contains(v * factor + u)) {74                nextCorrectGuesses++;75            }76            77            dfs2(v, u, nextCorrectGuesses);78        }79    }80}

Complexity

Time

O(N + M), where N is the number of nodes and M is the number of guesses. Building the adjacency list is O(N), populating the `guessSet` is O(M). Each of the two DFS passes visits every node and edge once, taking O(N) time. The final counting step takes O(N). Thus, the total time complexity is linear.

Space

O(N + M). O(N) for the adjacency list, O(M) for the `guessSet`, O(N) for the `correctGuessesCount` array, and O(N) for the recursion stack in the worst-case (for a skewed tree).

Trade-offs

Pros

Highly efficient, optimal solution that passes for large constraints.

Cons

More complex to conceptualize and implement compared to the brute-force approach.
Requires understanding of tree traversal and dynamic programming on trees.

Solutions

Solution

1class Solution {2private3  List<Integer>[] g;4private5  Map<Long, Integer> gs = new HashMap<>();6private7  int ans;8private9  int k;10private11  int cnt;12private13  int n;14public15  int rootCount(int[][] edges, int[][] guesses, int k) {16    this.k = k;17    n = edges.length + 1;18    g = new List[n];19    Arrays.setAll(g, e->new ArrayList<>());20    for (var e : edges) {21      int a = e[0], b = e[1];22      g[a].add(b);23      g[b].add(a);24    }25    for (var e : guesses) {26      int a = e[0], b = e[1];27      gs.merge(f(a, b), 1, Integer : : sum);28    }29    dfs1(0, -1);30    dfs2(0, -1);31    return ans;32  }33private34  void dfs1(int i, int fa) {35    for (int j : g[i]) {36      if (j != fa) {37        cnt += gs.getOrDefault(f(i, j), 0);38        dfs1(j, i);39      }40    }41  }42private43  void dfs2(int i, int fa) {44    ans += cnt >= k ? 1 : 0;45    for (int j : g[i]) {46      if (j != fa) {47        int a = gs.getOrDefault(f(i, j), 0);48        int b = gs.getOrDefault(f(j, i), 0);49        cnt -= a;50        cnt += b;51        dfs2(j, i);52        cnt -= b;53        cnt += a;54      }55    }56  }57private58  long f(int i, int j) { return 1L * i * n + j; }59}

1class Solution {2public:3  int rootCount(vector<vector<int>> &edges, vector<vector<int>> &guesses,4                int k) {5    int n = edges.size() + 1;6    vector<vector<int>> g(n);7    unordered_map<long long, int> gs;8    auto f = [&](int i, int j) { return 1LL * i * n + j; };9    for (auto &e : edges) {10      int a = e[0], b = e[1];11      g[a].push_back(b);12      g[b].push_back(a);13    }14    for (auto &e : guesses) {15      int a = e[0], b = e[1];16      gs[f(a, b)]++;17    }18    int ans = 0;19    int cnt = 0;20    function<void(int, int)> dfs1 = [&](int i, int fa) {21      for (int &j : g[i]) {22        if (j != fa) {23          cnt += gs[f(i, j)];24          dfs1(j, i);25        }26      }27    };28    function<void(int, int)> dfs2 = [&](int i, int fa) {29      ans += cnt >= k;30      for (int &j : g[i]) {31        if (j != fa) {32          int a = gs[f(i, j)];33          int b = gs[f(j, i)];34          cnt -= a;35          cnt += b;36          dfs2(j, i);37          cnt -= b;38          cnt += a;39        }40      }41    };42    dfs1(0, -1);43    dfs2(0, -1);44    return ans;45  }46};

1class Solution:2    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int: def dfs1(i, fa): nonlocal cnt for j in g[i]: if j != fa: cnt += gs[(i, j)] dfs1(j, i) def dfs2(i, fa): nonlocal ans, cnt ans += cnt >= k for j in g[i]: if j != fa: cnt -= gs[(i, j)] cnt += gs[(j, i)] dfs2(j, i) cnt -= gs[(j, i)] cnt += gs[(i, j)] g = defaultdict(list) for a, b in edges: g[a]. append(b) g[b]. append(a) gs = Counter((u, v) for u, v in guesses) cnt = 0 dfs1(0, - 1) ans = 0 dfs2(0, - 1) return ans

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem