0

Pick the root that needs the fewest edge reversals

by · 1 day ago· ⑂ 0 remixes

Question: Pick the root that needs the fewest edge reversals

Companies: Uber

Difficulty: Medium

Problem statement

You are given a tree with n nodes labeled 1 to n. Its n - 1 edges are directed: each edge (u, v) points from u to v. Ignoring direction, the edges connect all nodes into a tree.

You may pick any node as the root. Once the root is picked, every edge must point away from the root (from the endpoint closer to the root toward the one farther away). To get there you may reverse edges; each reversal costs 1.

Implement:

def min_root_reversals(n: int, edges: list[tuple[int, int]]) -> int

Return the minimum number of reversals over all choices of root.

Constraints

  • 1 <= n <= 10^5
  • len(edges) == n - 1, 1 <= u, v <= n, u != v
  • At most one edge between any pair of nodes; the undirected edges form a tree.

Example 1

n = 4, edges = [(2, 1), (2, 3), (4, 3)]
-> 1

Rooted at 2, edges 2 -> 1 and 2 -> 3 already point away; 4 -> 3 must become 3 -> 4, so 1 reversal. Rooting at 4 also costs 1 (flip 2 -> 3). No root costs 0.

Example 2

n = 4, edges = [(2, 1), (3, 1), (4, 1)]
-> 2

Every edge points into node 1. Rooting at a leaf, say 2, keeps 2 -> 1 and flips the other two.

Aim for better than trying every root with a separate traversal. Be ready to return the cost for every root as well.

Approach

Ignore direction for a moment: the n - 1 edges form a tree, so for any
chosen root there is exactly one way to orient every edge "away from the
root". An edge already points that way or it must be flipped. The cost of a
root is the number of edges on the wrong side.

Store each edge twice in an undirected adjacency list, with a weight:

  • u -> v as given: the forward entry (v, 0) costs nothing to traverse
    from u, the backward entry (u, 1) costs one reversal to traverse from v.

The cost of root r is then the sum of weights on all edges when you walk
the tree outward from r.

Brute force: try every root. Run a BFS/DFS from each node and sum the
weights. Each run is O(n), so O(n^2) total. Fine for n of a few thousand,
too slow for n = 10^5. Say it first, then improve it.

Optimal: reroot in two passes (O(n)).

  1. Root the tree at node 1. One traversal computes cost[1], the number of
    edges that point toward node 1's side.
  2. Move the root across one edge from parent p to child c. Only that
    edge changes its required direction; every other edge keeps the same
    "away from root" orientation, because the rest of the tree is still on
    the same side of it. So:
    • if the original edge is p -> c (weight 0 from p), it was correct
      and now must flip: cost[c] = cost[p] + 1;
    • if it is c -> p (weight 1 from p), it was flipped and now is
      correct: cost[c] = cost[p] - 1.
      Equivalently cost[c] = cost[p] + 1 - 2 * w(p, c).
  3. Answer is min(cost).

Invariants:

  • After pass 1, cost[1] is exact.
  • Each node's cost is set exactly once, from its parent's already-correct
    cost, in BFS order.

Use iterative BFS, not recursion. A path of 10^5 nodes overflows Python's
recursion limit, and interviewers at Uber have asked about this directly.

Code

from collections import deque
from typing import List, Tuple


def min_reversals_all_roots(n: int, edges: List[Tuple[int, int]]) -> List[int]:
    """cost[r] = reversals needed so every edge points away from root r.

    Nodes are labeled 1..n; cost[0] is unused (set to 0).
    """
    adj: List[List[Tuple[int, int]]] = [[] for _ in range(n + 1)]
    for u, v in edges:
        adj[u].append((v, 0))  # walking u -> v follows the edge
        adj[v].append((u, 1))  # walking v -> u needs one reversal

    # Pass 1: BFS from node 1, record order and parent-edge weights.
    parent = [0] * (n + 1)
    weight_from_parent = [0] * (n + 1)
    order: List[int] = []
    seen = [False] * (n + 1)
    seen[1] = True
    queue = deque([1])
    base = 0
    while queue:
        node = queue.popleft()
        order.append(node)
        for nxt, w in adj[node]:
            if not seen[nxt]:
                seen[nxt] = True
                parent[nxt] = node
                weight_from_parent[nxt] = w
                base += w
                queue.append(nxt)

    if len(order) != n:
        raise ValueError("edges do not form a connected tree")

    # Pass 2: reroot along BFS order; a parent is always finished first.
    cost = [0] * (n + 1)
    cost[1] = base
    for node in order[1:]:
        w = weight_from_parent[node]
        cost[node] = cost[parent[node]] + 1 - 2 * w
    return cost


def min_root_reversals(n: int, edges: List[Tuple[int, int]]) -> int:
    if n == 1:
        return 0
    return min(min_reversals_all_roots(n, edges)[1:])


def brute_force(n: int, edges: List[Tuple[int, int]]) -> int:
    adj: List[List[Tuple[int, int]]] = [[] for _ in range(n + 1)]
    for u, v in edges:
        adj[u].append((v, 0))
        adj[v].append((u, 1))
    best = n
    for root in range(1, n + 1):
        total, seen, stack = 0, {root}, [root]
        while stack:
            node = stack.pop()
            for nxt, w in adj[node]:
                if nxt not in seen:
                    seen.add(nxt)
                    total += w
                    stack.append(nxt)
        best = min(best, total)
    return best


if __name__ == "__main__":
    import random

    assert min_root_reversals(1, []) == 0
    assert min_root_reversals(2, [(2, 1)]) == 0
    assert min_root_reversals(4, [(2, 1), (2, 3), (4, 3)]) == 1
    assert min_root_reversals(4, [(1, 2), (1, 3), (1, 4)]) == 0
    assert min_root_reversals(4, [(2, 1), (3, 1), (4, 1)]) == 2
    assert min_reversals_all_roots(3, [(1, 2), (3, 2)])[1:] == [1, 2, 1]

    # Long path: no recursion, so no stack overflow.
    n = 100_000
    assert min_root_reversals(n, [(i + 1, i) for i in range(1, n)]) == 0

    for _ in range(300):
        n = random.randint(1, 9)
        edges = []
        for v in range(2, n + 1):
            u = random.randint(1, v - 1)
            edges.append((u, v) if random.random() < 0.5 else (v, u))
        assert min_root_reversals(n, edges) == brute_force(n, edges)
    print("ok")

Complexity

  • Building the adjacency list: O(n) time, O(n) space.
  • Pass 1 (BFS from node 1): O(n) time.
  • Pass 2 (rerooting): O(n) time, O(1) work per node.
  • Total: O(n) time, O(n) space. The brute force is O(n^2) time.

Test cases

  • n = 1, edges = [] -> 0 (a single node needs nothing).
  • n = 2, edges = [(2, 1)] -> 0 (root at 2).
  • n = 4, edges = [(2, 1), (2, 3), (4, 3)] -> 1. Root 2 needs 4 -> 3
    flipped; root 4 needs 2 -> 3 flipped. No root needs 0.
  • Out-star n = 4, edges = [(1, 2), (1, 3), (1, 4)] -> 0 (root 1).
  • In-star n = 4, edges = [(2, 1), (3, 1), (4, 1)] -> 2 (root at any leaf,
    flip the other two leaves' edges).
  • Per-root costs for n = 3, edges = [(1, 2), (3, 2)] -> [1, 2, 1].
  • Path of 100,000 nodes all pointing one way -> 0, and no recursion error.
  • Random trees up to 9 nodes: matches the brute force.

Follow-ups

  • Return the cost for every root (LeetCode 2858 style). The cost
    array already holds it; return it instead of the minimum.
  • Return all roots that achieve the minimum (the Codeforces "choosing a
    capital" variant). Scan cost once and collect indices equal to the min.
  • Weighted reversals (flipping edge e costs c_e). Same rerooting:
    store the flip cost on the backward entry, and moving the root across an
    edge changes the total by +c_e (edge was correct) or -c_e (edge was
    flipped). Still O(n).
  • Input is a forest or may contain a cycle. Validate: exactly n - 1
    edges and BFS reaches all n nodes, else raise or return -1. For a
    general directed graph, "every node reachable from the root with minimum
    reversals" becomes a 0-1 BFS per root; for a forest, solve each component.
  • Edges arrive online and queries ask the current best root. Discuss
    that one edge flip changes the cost of every node on one side of that edge
    by ±1, so it is a subtree range update on an Euler tour; a segment tree
    with range add and range min answers it in O(log n).

Interviewer expectations

Mid-level (L4)

  • Sees that the cost of a root is "edges pointing toward it" and writes the
    O(n^2) brute force correctly.
  • With a hint about moving the root one step, derives the ±1 update.
  • Handles n = 1 and 1-based labels.

Senior (L5)

  • Reaches the two-pass rerooting on their own and explains why only the
    crossed edge changes.
  • Writes it iteratively and calls out the recursion limit on a 10^5 path.
  • Tests with a star, a path and a small hand-checked case, and can return
    all costs or all optimal roots without rewriting.

Staff (L6+)

  • Frames rerooting as a general tree-DP technique and applies it to the
    weighted version immediately.
  • Validates input (not a tree, disconnected) and states the behavior.
  • Reasons about the dynamic follow-up (Euler tour plus range-add/min
    segment tree) and the trade-offs against recomputing.

Sign in to comment, vote or remix.