Count number of possible root nodes

Array Hash Table Dynamic Programming Tree Depth-First Search

48.1% acceptance

Feb 25, 2026

322

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] = [ai, bi] indicates that there is an edge between nodes ai and bi 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] = [uj, vj] indicates Bob guessed uj to be the parent of vj. 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.

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

impl Solution {
  pub fn root_count(edges: Vec<Vec<i32>>, guesses: Vec<Vec<i32>>, k: i32) -> i32 {
    use std::collections::HashSet;
    let n = edges.len() + 1;
    let mut adj: Vec<Vec<usize>> = vec![vec![]; n];
    for e in &edges {
      let (a, b) = (e[0] as usize, e[1] as usize);
      adj[a].push(b);
      adj[b].push(a);
    }
    let guess_set: HashSet<(usize, usize)> = guesses.iter()
      .map(|g| (g[0] as usize, g[1] as usize))
      .collect();

    // DFS from root 0 to compute initial score
    let mut score0 = 0i32;
    let mut parent = vec![usize::MAX; n];
    let mut order = Vec::with_capacity(n);
    let mut stack = vec![0usize];
    parent[0] = 0;
    while let Some(u) = stack.pop() {
      order.push(u);
      for &v in &adj[u] {
        if parent[v] == usize::MAX && v != 0 {
          parent[v] = u;
          if guess_set.contains(&(u, v)) { score0 += 1; }
          stack.push(v);
        }
      }
    }

    // Rerooting DP
    let mut scores = vec![0i32; n];
    scores[0] = score0;
    let mut ans = 0i32;
    for &u in order.iter() {
      if scores[u] >= k { ans += 1; }
      for &v in &adj[u] {
        if parent[v] == u {
          let mut s = scores[u];
          if guess_set.contains(&(u, v)) { s -= 1; }
          if guess_set.contains(&(v, u)) { s += 1; }
          scores[v] = s;
        }
      }
    }
    ans
  }
}