0

Minimum distance between two nodes in an N-ary tree

by Frontier Interview· 1 day ago· ⑂ 0 remixes

Question: Minimum distance between two nodes in an N-ary tree

Companies: Scale AI, Google

Difficulty: Medium

Problem statement

You are given the root of an N-ary tree. Every node has an integer val, unique within the tree, and a list children of its direct children. There are no parent pointers.

Implement:

def min_distance(root: Node | None, u: int, v: int) -> int

Return the number of edges on the path between the node whose value is u and the node whose value is v. If either value is not in the tree, return -1. If u == v and the value exists, return 0.

Constraints

  • 0 <= number of nodes <= 10^5
  • The tree can be very deep (for example, a single chain), so do not rely on deep recursion.
  • One node may be an ancestor of the other.

Example 1

        10
     /   |   \
   20    30   40
  /  \         |
 50   60       70
      |
      80
  • min_distance(root, 50, 80) returns 3: the path is 50 → 20 → 60 → 80.
  • min_distance(root, 80, 70) returns 5: the path goes up to 10 and down to 70.

Example 2

Same tree. min_distance(root, 20, 80) returns 2, because 20 is an ancestor of 80. min_distance(root, 30, 99) returns -1, because 99 is not in the tree.

Be ready to discuss how you would answer many queries on the same tree quickly.

Approach

The distance between two nodes is the number of edges on the path that joins them. In a tree that path is unique. It climbs from u to the lowest common ancestor (LCA) of u and v, then descends to v. So:

dist(u, v) = depth(u) + depth(v) - 2 * depth(lca(u, v))

The tree only gives child pointers, so you cannot walk upward. The work is to recover enough parent information, or to answer in one downward pass.

Brute force. Find the root-to-u path and the root-to-v path with two DFS calls. Walk both paths from the root while they agree. The last shared node is the LCA. The answer is len(path_u) + len(path_v) - 2 * len(shared). This is O(n) time and O(h) extra space, and it is already correct. It is a good first answer. Its cost is two traversals plus two path copies.

One pass (the LCA 236 pattern, carrying distances). Run a post-order DFS. Each subtree returns a pair: the distance from its root down to u if u is inside, else None, and the same for v. At a node, take the pair from each child and add 1 to each found distance. If the node itself is u or v, its distance to that target is 0. The first node where both distances are known is the LCA. Its answer is du + dv. Return it immediately and stop the search.

This also covers the case where one target is an ancestor of the other. At that node, one distance is 0 and the other comes from a child. When u == v, the answer is 0 at that node.

If either target is missing from the tree, the root never sees both distances. Return -1. LeetCode 236 assumes both nodes exist. Say out loud that you do not assume that here.

The DFS is iterative with an explicit stack. A path-shaped tree with 10^5 nodes would overflow Python's recursion limit.

Many queries. For q queries on the same tree, run one BFS to build parent and depth for every value. Then build a binary-lifting table: up[k][x] is the 2^k-th ancestor of x. A query lifts the deeper node to the same depth, then lifts both nodes together while their ancestors differ. Each query costs O(log n).

Invariants.

  • In the post-order pass, a node's pair is final only after all of its children are processed.
  • The first node with both distances is the LCA. Any lower node lacks one of the targets, and any higher node is processed later.
  • In binary lifting, depth[up[k][x]] == depth[x] - 2^k whenever that ancestor exists. The root's ancestor is itself.

Code

from collections import deque
from typing import Dict, List, Optional, Tuple


class Node:
    def __init__(self, val: int, children: Optional[List["Node"]] = None) -> None:
        self.val = val
        self.children: List["Node"] = children or []


def min_distance(root: Optional[Node], u: int, v: int) -> int:
    """Edges between the nodes with values u and v, or -1 if either is missing."""
    if root is None:
        return -1
    # Each stack entry: (node, children_done). Results map a node id to (du, dv).
    found: Dict[int, Tuple[Optional[int], Optional[int]]] = {}
    stack: List[Tuple[Node, bool]] = [(root, False)]
    while stack:
        node, done = stack.pop()
        if not done:
            stack.append((node, True))
            for child in node.children:
                stack.append((child, False))
            continue
        du = 0 if node.val == u else None
        dv = 0 if node.val == v else None
        for child in node.children:
            cu, cv = found.pop(id(child))
            if du is None and cu is not None:
                du = cu + 1
            if dv is None and cv is not None:
                dv = cv + 1
        if du is not None and dv is not None:
            return du + dv  # this node is the LCA
        found[id(node)] = (du, dv)
    return -1


class TreeDistance:
    """Preprocess once, then answer each distance query in O(log n)."""

    def __init__(self, root: Node) -> None:
        self.depth: Dict[int, int] = {root.val: 0}
        parent: Dict[int, int] = {root.val: root.val}
        queue = deque([root])
        while queue:
            node = queue.popleft()
            for child in node.children:
                parent[child.val] = node.val
                self.depth[child.val] = self.depth[node.val] + 1
                queue.append(child)
        levels = max(1, max(self.depth.values()).bit_length())
        self.up: List[Dict[int, int]] = [parent]
        for _ in range(1, levels):
            prev = self.up[-1]
            self.up.append({x: prev[prev[x]] for x in prev})

    def lca(self, a: int, b: int) -> int:
        if self.depth[a] < self.depth[b]:
            a, b = b, a
        diff = self.depth[a] - self.depth[b]
        k = 0
        while diff:
            if diff & 1:
                a = self.up[k][a]
            diff >>= 1
            k += 1
        if a == b:
            return a
        for k in range(len(self.up) - 1, -1, -1):
            if self.up[k][a] != self.up[k][b]:
                a, b = self.up[k][a], self.up[k][b]
        return self.up[0][a]

    def distance(self, a: int, b: int) -> int:
        if a not in self.depth or b not in self.depth:
            return -1
        return self.depth[a] + self.depth[b] - 2 * self.depth[self.lca(a, b)]

Complexity

  • min_distance: O(n) time. O(n) extra space in the worst case for the stack and the pending results. It stops early once it reaches the LCA.
  • TreeDistance.__init__: O(n log h) time and space, where h is the tree height.
  • TreeDistance.lca and distance: O(log h) time, O(1) extra space.
  • Brute force (two root paths): O(n) time, O(h) extra space per query.

Test cases

Tree used below: 1 -> [2, 3, 4], 2 -> [5, 6], 4 -> [7], 6 -> [8].

  • min_distance(root, 5, 8) → 3 (5-2-6-8; LCA is 2).
  • min_distance(root, 8, 7) → 5 (8-6-2-1-4-7; LCA is the root).
  • min_distance(root, 2, 8) → 2 (2 is an ancestor of 8).
  • min_distance(root, 3, 3) → 0 (same node).
  • min_distance(root, 1, 7) → 2 (one node is the root).
  • min_distance(root, 5, 99) → -1 (99 is not in the tree).
  • min_distance(None, 1, 2) → -1 (empty tree).
  • Single node 1: min_distance(root, 1, 1) → 0.
  • Path of 100,000 nodes 0 -> 1 -> ... -> 99999: min_distance(root, 0, 99999) → 99999 with no recursion error.
  • TreeDistance(root).distance(a, b) matches min_distance for every pair in the sample tree.

Follow-ups

  1. Many queries on a static tree. Preprocess with BFS plus binary lifting (TreeDistance). Each query is O(log n). An Euler tour plus a sparse table gives O(1) LCA queries after O(n log n) preprocessing.
  2. Nodes may be missing, or values may repeat. Missing nodes return -1, as above. With duplicate values, the question is ill-defined. Ask whether to use node references instead of values, or to return the minimum over all matching pairs. The minimum over all pairs needs a multi-source BFS on the undirected tree.
  3. Distance from one node to a set of nodes, or the closest marked node. Convert the tree to an undirected adjacency list (child edges plus parent edges). Run a BFS from the start node, or a multi-source BFS from all marked nodes. This is O(n).
  4. Weighted edges. Store a weighted depth (distance from the root) next to the edge-count depth. Use the edge-count depth for lifting, and return wdepth[u] + wdepth[v] - 2 * wdepth[lca].
  5. The tree changes (insert a leaf, then query). A new leaf's lifting row depends only on its ancestors, so it is O(log n) to add. Deleting subtrees or re-parenting needs link-cut trees, or rebuilding when changes are rare. Discuss the tradeoff instead of coding it.

Interviewer expectations

Mid-level (L4)

  • States the depth formula or the root-path idea and gets a correct O(n) solution.
  • Handles the cases where one node is an ancestor of the other and where u == v.
  • Walks through the sample tree by hand and gives the correct complexity.

Senior (L5)

  • Writes the single-pass DFS that returns distances, not only the LCA.
  • Asks whether both nodes are guaranteed to exist, and returns -1 if not.
  • Notices the recursion depth risk on a skewed tree and uses an explicit stack.
  • Proposes preprocessing for repeated queries without being prompted.

Staff (L6+)

  • Implements binary lifting cleanly and compares it with Euler tour plus sparse table (preprocessing cost, memory, query time).
  • Reasons about duplicate values, node references versus values, and weighted edges.
  • Discusses dynamic trees, and when rebuilding is cheaper than a complex data structure.

Sign in to comment, vote or remix.