2581. Count Number of Possible Root Nodes

• Time: $O(n)$
• Space: $O(n)$

C++JavaPython

class Solution {
 public:
  int rootCount(vector<vector<int>>& edges, vector<vector<int>>& guesses,
                int k) {
    int ans = 0;
    const int n = edges.size() + 1;
    vector<vector<int>> graph(n);
    vector<unordered_set<int>> guessGraph(n);
    vector<int> parent(n);

    for (const vector<int>& edge : edges) {
      const int u = edge[0];
      const int v = edge[1];
      graph[u].push_back(v);
      graph[v].push_back(u);
    }

    for (const vector<int>& guess : guesses) {
      const int u = guess[0];
      const int v = guess[1];
      guessGraph[u].insert(v);
    }

    // Precalculate `parent`.
    dfs(graph, 0, /*prev=*/-1, parent);

    // Calculate `correctGuess` for tree rooted at 0.
    int correctGuess = 0;
    for (int i = 1; i < n; ++i)
      if (guessGraph[parent[i]].contains(i))
        ++correctGuess;

    reroot(graph, 0, /*prev=*/-1, correctGuess, guessGraph, k, ans);
    return ans;
  }

 private:
  void dfs(const vector<vector<int>>& graph, int u, int prev,
           vector<int>& parent) {
    parent[u] = prev;
    for (const int v : graph[u])
      if (v != prev)
        dfs(graph, v, u, parent);
  }

  void reroot(const vector<vector<int>>& graph, int u, int prev,
              int correctGuess, const vector<unordered_set<int>>& guessGraph,
              const int& k, int& ans) {
    if (u != 0) {
      // The tree is rooted at u, so a guess edge (u, prev) will match the new
      // `parent` relationship.
      if (guessGraph[u].contains(prev))
        ++correctGuess;
      // A guess edge (prev, u) matching the old `parent` relationship will no
      // longer be true.
      if (guessGraph[prev].contains(u))
        --correctGuess;
    }
    if (correctGuess >= k)
      ++ans;
    for (const int v : graph[u])
      if (v != prev)
        reroot(graph, v, u, correctGuess, guessGraph, k, ans);
  }
};

class Solution {
  public int rootCount(int[][] edges, int[][] guesses, int k) {
    final int n = edges.length + 1;
    List<Integer>[] graph = new List[n];
    Set<Integer>[] guessGraph = new Set[n];
    int[] parent = new int[n];

    for (int i = 0; i < n; ++i) {
      graph[i] = new ArrayList<>();
      guessGraph[i] = new HashSet<>();
    }

    for (int[] edge : edges) {
      final int u = edge[0];
      final int v = edge[1];
      graph[u].add(v);
      graph[v].add(u);
    }

    for (int[] guess : guesses) {
      final int u = guess[0];
      final int v = guess[1];
      guessGraph[u].add(v);
    }

    // Precalculate `parent`.
    dfs(graph, 0, /*prev=*/-1, parent);

    // Calculate `correctGuess` for tree rooted at 0.
    int correctGuess = 0;
    for (int i = 1; i < n; ++i)
      if (guessGraph[parent[i]].contains(i))
        ++correctGuess;

    reroot(graph, 0, /*prev=*/-1, correctGuess, guessGraph, k);
    return ans;
  }

  private int ans = 0;

  private void dfs(List<Integer>[] graph, int u, int prev, int[] parent) {
    parent[u] = prev;
    for (final int v : graph[u])
      if (v != prev)
        dfs(graph, v, u, parent);
  }

  private void reroot(List<Integer>[] graph, int u, int prev, int correctGuess,
                      Set<Integer>[] guessGraph, int k) {
    if (u != 0) {
      // The tree is rooted at u, so a guess edge (u, prev) will match the new
      // `parent` relationship.
      if (guessGraph[u].contains(prev))
        ++correctGuess;
      // A guess edge (prev, u) matching the old `parent` relationship will no
      // longer be true.
      if (guessGraph[prev].contains(u))
        --correctGuess;
    }
    if (correctGuess >= k)
      ++ans;
    for (final int v : graph[u])
      if (v != prev)
        reroot(graph, v, u, correctGuess, guessGraph, k);
  }
}

class Solution:
  def rootCount(
      self,
      edges: list[list[int]],
      guesses: list[list[int]],
      k: int,
  ) -> int:
    ans = 0
    n = len(edges) + 1
    graph = [[] for _ in range(n)]
    guessGraph = [set() for _ in range(n)]
    parent = [0] * n

    for u, v in edges:
      graph[u].append(v)
      graph[v].append(u)

    for u, v in guesses:
      guessGraph[u].add(v)

    def dfs(u: int, prev: int) -> None:
      parent[u] = prev
      for v in graph[u]:
        if v != prev:
          dfs(v, u)

    # Precalculate `parent`.
    dfs(0, -1)

    # Calculate `correctGuess` for tree rooted at 0.
    correctGuess = sum(i in guessGraph[parent[i]] for i in range(1, n))

    def reroot(u: int, prev: int, correctGuess: int) -> None:
      nonlocal ans
      if u != 0:
        # The tree is rooted at u, so a guess edge (u, prev) will match the new
        # `parent` relationship.
        if prev in guessGraph[u]:
          correctGuess += 1
        # A guess edge (prev, u) matching the old `parent` relationship will no
        # longer be True.
        if u in guessGraph[prev]:
          correctGuess -= 1
      if correctGuess >= k:
        ans += 1
      for v in graph[u]:
        if v != prev:
          reroot(v, u, correctGuess)

    reroot(0, -1, correctGuess)
    return ans