Root of the largest tree in a forest: Frontier Interview's solution

Question: Root of the largest tree in a forest

Companies: Google, Microsoft

Difficulty: Medium

Problem statement

You are given a list relations. Each entry [child, parent] says that parent is the direct parent of child. Together the relations form a forest: there are no cycles, and every node has at most one parent. Node IDs are integers, and the forest contains only the nodes that appear in relations.

A root is a node that appears in relations but never as a child. The size of a tree is the number of distinct nodes in it, root included.

Implement:

def largest_tree_root(relations: list[list[int]]) -> int | None

The function returns the root of the tree with the largest size. If two or more trees have the same size, return the smallest root ID. If relations is empty, return None.

Constraints

0 <= len(relations) <= 10^5
Node IDs are non-negative integers up to 10^9.
A tree can be a single chain up to 10^5 nodes deep.

Example 1

relations = [[4, 9], [7, 9], [1, 7], [3, 2]]

The tree rooted at 9 holds {9, 4, 7, 1}, so its size is 4. The tree rooted at 2 holds {2, 3}, so its size is 2. Answer: 9.

Example 2

relations = [[15, 20], [16, 20], [31, 30], [32, 30]]

Both trees have 3 nodes, so the smaller root wins. Answer: 20.

Reconstructed from the reported question; details may differ from what interviewers ask.

Approach

Every pair [child, parent] is an edge from a parent down to a child. Because
each node has at most one parent and there are no cycles, the nodes split into
separate trees. A root is a node that shows up as a parent but never as a child.
We need the size of each tree, and then the root of the biggest one. On a tie
we return the smallest root ID.

Brute force. For every node, follow parent pointers up until you reach a
node with no parent, then add one to that root's count. This is correct, but a
long chain costs O(n) per node, so the total is O(n^2) in the worst case.

Optimal: count down from each root. Build a children map and a set of
every node that appears as a child. The roots are all nodes minus that set. Run
an iterative DFS (an explicit stack) from each root and count the nodes it
reaches. Each node belongs to exactly one tree, so the walks visit every node
once in total. That is O(n).

Alternative: memoised walk up. Keep root_of[node]. Walk up from a node
until you hit a node whose root is already known, then write the root into
every node on the path. This is path compression without union-find, and it is
also O(n). It works well when the input only gives parent pointers.

Invariants:

Every node appears in exactly one tree, so the tree sizes add up to the
number of distinct nodes.
A node is a root exactly when it is not in the is_child set.
Tie rule: compare (size, -root), so a larger size wins and an equal size
goes to the smaller ID.

Ask the interviewer about the tie rule and the empty input early. Those are the
two details the reported prompts disagree on.

Code

from collections import defaultdict
from typing import Optional


def largest_tree_root(relations: list[list[int]]) -> Optional[int]:
    """Return the root of the largest tree in the forest.

    relations: [child, parent] pairs forming an acyclic forest where each
    child has at most one parent. Ties go to the smallest root id.
    Returns None for an empty forest.
    """
    children: dict[int, list[int]] = defaultdict(list)
    nodes: set[int] = set()
    is_child: set[int] = set()
    for child, parent in relations:
        children[parent].append(child)
        nodes.add(child)
        nodes.add(parent)
        is_child.add(child)

    best_root: Optional[int] = None
    best_size = 0
    for root in sorted(nodes - is_child):
        size = 0
        stack = [root]
        while stack:
            node = stack.pop()
            size += 1
            stack.extend(children.get(node, ()))
        # Roots are visited in ascending order, so strict > keeps the smallest id on ties.
        if size > best_size:
            best_root, best_size = root, size
    return best_root


def tree_sizes(relations: list[list[int]]) -> dict[int, int]:
    """Follow-up helper: root -> size, using memoised walks up the parent pointers."""
    parent_of = {c: p for c, p in relations}
    nodes = set(parent_of) | set(parent_of.values())
    root_of: dict[int, int] = {}
    for start in nodes:
        path = []
        node = start
        while node not in root_of and node in parent_of:
            path.append(node)
            node = parent_of[node]
        root = root_of.get(node, node)
        root_of[node] = root
        for n in path:
            root_of[n] = root
    sizes: dict[int, int] = defaultdict(int)
    for n in nodes:
        sizes[root_of[n]] += 1
    return dict(sizes)


if __name__ == "__main__":
    assert largest_tree_root([]) is None
    assert largest_tree_root([[2, 1]]) == 1
    # Tree 10: {10, 11, 12, 13}; tree 5: {5, 6}
    rel = [[11, 10], [12, 10], [13, 12], [6, 5]]
    assert largest_tree_root(rel) == 10
    assert tree_sizes(rel) == {10: 4, 5: 2}
    # Tie: both trees have 3 nodes -> smaller root wins.
    assert largest_tree_root([[8, 7], [9, 7], [2, 3], [4, 3]]) == 3
    # Long chain 0 <- 1 <- ... <- 9999 (no recursion limit issues).
    chain = [[i + 1, i] for i in range(9999)]
    assert largest_tree_root(chain + [[20001, 20000]]) == 0
    assert tree_sizes(chain)[0] == 10000
    print("ok")

Complexity

Let n be the number of relations. The forest has at most n + 1 distinct
nodes per tree, and O(n) nodes overall.

Building the maps: O(n) time and O(n) space.
Finding roots: O(n) to take the set difference. Sorting the roots costs
O(r log r) for r roots. You can drop the sort by comparing
(size, -root) directly, which makes the whole thing O(n).
DFS over all trees: O(n) time in total, because each node is pushed once.
The stack holds at most O(n) entries.
tree_sizes: O(n) amortised, because each node's root is written once and
never walked again.

Test cases

[] -> None. The forest is empty. Confirm with the interviewer: some
versions expect -1.
[[2, 1]] -> 1. One tree with two nodes.
[[11, 10], [12, 10], [13, 12], [6, 5]] -> 10. The tree sizes are 4
and 2.
[[8, 7], [9, 7], [2, 3], [4, 3]] -> 3. Both trees have 3 nodes, so the
smaller root wins.
A chain 0 <- 1 <- ... <- 9999 plus a separate pair 20001 <- 20000 ->
0. This checks that the DFS is iterative and does not hit Python's
recursion limit.
Many single-edge trees [[2k+1, 2k]]: every tree has size 2, so the
answer is the smallest root, 0.
A star where one parent has 10^5 children -> that parent.

Follow-ups

Relations arrive as a stream, and you are asked for the current answer
after each one. Use union-find with size. When [c, p] arrives, union
the set of c under the set of p. Keep the representative pointing at
the real root by always attaching c's set under p's set. Union by size
no longer works because the direction is fixed, so use path compression
for amortised near-O(log n). Keep a running best (size, -root). Sizes
only grow, so the best value never has to be taken back.
Return the size of the tree that contains any node x, for many
queries. Precompute root_of with the memoised walk up (tree_sizes),
then answer each query in O(1).
The input may be invalid (cycles, or a child with two parents).
Detect a second parent while building parent_of. To find a cycle, look
for nodes the DFS from roots never visits: any node left unvisited sits on
a cycle or hangs below one. Report it, or raise an error.
Weight the trees. Each node has a value, and the "largest" tree is the
one with the largest sum, or the deepest one. Use the same DFS, but track
the sum, or carry (node, depth) on the stack. The tie rule stays the
same.
The relations do not fit in memory (billions of edges). Run the
walk-up in rounds as a MapReduce job: pointer jumping, where in each round
parent = parent[parent]. After O(log depth) rounds every node points at
its root. Then group by root and count.

Interviewer expectations

Mid-level (L4)

Works out that roots are the parents that never appear as a child, and
builds the child adjacency list.
Writes a correct DFS or BFS that counts each tree, and handles the tie rule
once it is pointed out.
States O(n) time and O(n) space.

Senior (L5)

Asks about the tie rule, the empty input, and whether the input is
guaranteed valid, before writing code.
Uses an iterative traversal, or explains the recursion-depth risk on long
chains.
Starts from the O(n^2) walk-up and improves it with memoisation without
being prompted. Writes clean tests, including edge cases.

Staff (L6+)

Handles the streaming version with union-find, and explains why union by
size conflicts with keeping the real root as the representative.
Discusses how to validate input (cycles, nodes with two parents) and what to
return in those cases.
Outlines a distributed approach (pointer jumping or connected components)
for a forest that does not fit on one machine.