1719. Number Of Ways To Reconstruct A Tree

• Time: $O(n)$
• Space: $O(500 + n) = O(n)$

C++JavaPython

class Solution {
 public:
  int checkWays(vector<vector<int>>& pairs) {
    constexpr int kMax = 501;
    unordered_map<int, vector<int>> graph;
    vector<int> degrees(kMax);
    vector<vector<bool>> connected(kMax, vector<bool>(kMax));

    for (const vector<int>& pair : pairs) {
      const int u = pair[0];
      const int v = pair[1];
      graph[u].push_back(v);
      graph[v].push_back(u);
      ++degrees[u];
      ++degrees[v];
      connected[u][v] = true;
      connected[v][u] = true;
    }

    // For each node, sort its children by degrees in descending order.
    for (auto& [_, children] : graph)
      ranges::sort(children, ranges::greater{},
                   [&degrees](int child) { return degrees[child]; });

    const int root = getRoot(degrees, graph.size());
    if (root == -1)
      return 0;
    if (!dfs(graph, root, degrees, connected, {}, vector<bool>(kMax)))
      return 0;
    return hasMoreThanOneWay ? 2 : 1;
  }

 private:
  bool hasMoreThanOneWay = false;

  // Returns the root by finding the node with a degree that equals to n - 1.
  int getRoot(const vector<int>& degrees, int n) {
    for (int i = 1; i < degrees.size(); ++i)
      if (degrees[i] == n - 1)
        return i;
    return -1;
  }

  // Returns true if each node rooted at u is connected to all of its ancestors.
  bool dfs(const unordered_map<int, vector<int>>& graph, int u,
           vector<int>& degrees, vector<vector<bool>>& connected,
           vector<int>&& ancestors, vector<bool>&& seen) {
    seen[u] = true;

    for (const int ancestor : ancestors)
      if (!connected[u][ancestor])
        return false;

    ancestors.push_back(u);

    for (const int v : graph.at(u)) {
      if (seen[v])
        continue;
      // We can swap u with v, so there are more than one way.
      if (degrees[v] == degrees[u])
        hasMoreThanOneWay = true;
      if (!dfs(graph, v, degrees, connected, std::move(ancestors),
               std::move(seen)))
        return false;
    }

    ancestors.pop_back();
    return true;
  }
};

class Solution {
  public int checkWays(int[][] pairs) {
    final int kMax = 501;
    Map<Integer, List<Integer>> graph = new HashMap<>();
    int[] degrees = new int[kMax];
    boolean[][] connected = new boolean[kMax][kMax];

    for (int[] pair : pairs) {
      final int u = pair[0];
      final int v = pair[1];
      graph.putIfAbsent(u, new ArrayList<>());
      graph.putIfAbsent(v, new ArrayList<>());
      graph.get(u).add(v);
      graph.get(v).add(u);
      ++degrees[u];
      ++degrees[v];
      connected[u][v] = true;
      connected[v][u] = true;
    }

    // For each node, sort its children by degrees in descending order.
    for (final int u : graph.keySet())
      graph.get(u).sort((a, b) -> Integer.compare(degrees[b], degrees[a]));

    final int root = getRoot(degrees, graph.keySet().size());
    if (root == -1)
      return 0;
    if (!dfs(graph, root, degrees, connected, new ArrayList<>(), new boolean[kMax]))
      return 0;
    return hasMoreThanOneWay ? 2 : 1;
  }

  private boolean hasMoreThanOneWay = false;

  // Returns the root by finding the node with a degree that equals to n - 1.
  private int getRoot(int[] degrees, int n) {
    for (int i = 1; i < degrees.length; ++i)
      if (degrees[i] == n - 1)
        return i;
    return -1;
  }

  // Returns true if each node rooted at u is connected to all of its ancestors.
  private boolean dfs(Map<Integer, List<Integer>> graph, int u, int[] degrees,
                      boolean[][] connected, List<Integer> ancestors, boolean[] seen) {
    seen[u] = true;

    for (final int ancestor : ancestors)
      if (!connected[u][ancestor])
        return false;

    ancestors.add(u);

    for (final int v : graph.get(u)) {
      if (seen[v])
        continue;
      // We can swap u with v, so there are more than one way.
      if (degrees[v] == degrees[u])
        hasMoreThanOneWay = true;
      if (!dfs(graph, v, degrees, connected, ancestors, seen))
        return false;
    }

    ancestors.remove(ancestors.size() - 1);
    return true;
  }
}

class Solution:
  def checkWays(self, pairs: list[list[int]]) -> int:
    kMax = 501
    graph = collections.defaultdict(list)
    degrees = [0] * kMax
    connected = [[False] * kMax for _ in range(kMax)]

    for u, v in pairs:
      graph[u].append(v)
      graph[v].append(u)
      degrees[u] += 1
      degrees[v] += 1
      connected[u][v] = True
      connected[v][u] = True

    # For each node, sort its children by degrees in descending order.
    for _, children in graph.items():
      children.sort(key=lambda x: -degrees[x])

    # Find the root with a degree that equals to n - 1.
    root = next((i for i, d in enumerate(degrees) if d == len(graph) - 1), -1)
    if root == -1:
      return 0

    hasMoreThanOneWay = False

    def dfs(u: int, ancestors: list[int], seen: list[bool]) -> bool:
      """
      Returns True if each node rooted at u is connected to all of its
      ancestors.
      """
      nonlocal hasMoreThanOneWay
      seen[u] = True
      for ancestor in ancestors:
        if not connected[u][ancestor]:
          return False
      ancestors.append(u)
      for v in graph[u]:
        if seen[v]:
          continue
        if degrees[v] == degrees[u]:
          hasMoreThanOneWay = True
        if not dfs(v, ancestors, seen):
          return False
      ancestors.pop()
      return True

    if not dfs(root, [], [False] * kMax):
      return 0
    return 2 if hasMoreThanOneWay else 1