Compress an n-ary tree by promoting children

Clarify first. Values are unique, so a node can be named by its value. The root is never compressed. When a node is compressed, it disappears and its children take its place in its parent's child list, in their original order. Ask these questions before you write code:

• Can a compressed node have compressed children? (Yes. Then its grandchildren move up two levels.)
• Do promoted children go where the removed node was, or at the end of the parent's list? (Assume where it was. That keeps the left-to-right order of the remaining nodes.)
• Are values in the list that are not in the tree an error? (Assume they are ignored.)
• Can we change the tree in place? (Assume yes, and return the same root.)

Brute force. For each value in the list, search the tree for its parent, splice the node's children into the parent's list at the node's position, and remove the node. Each removal costs O(n) to find the parent, plus the cost of the list splice. With k values to compress, that is O(n * k). It works, and it is a fine thing to say out loud first.

Optimal: one post-order pass. Put the values in a set. Define kept(node) as the list of nodes that should stand in this node's place in its parent's child list:

• If node is kept, its new children are the concatenation of kept(child) for each child, in order. Then kept(node) = [node].
• If node is compressed, kept(node) is that same concatenation, and node itself is gone.

Process children before their parent (post-order), so every child's kept list is ready when the parent needs it. A chain of compressed nodes collapses naturally, because a compressed child already returns its own kept descendants.

Invariant. After a node is processed, its subtree holds no compressed node below it, and the kept nodes appear in the same pre-order as in the original tree. The root is always kept, so the answer is the root itself.

Why iterative. A tree of 10^5 nodes can be a single path, and Python's default recursion limit is 1000. The code uses an explicit stack. Say this in the interview, then show the recursive version in two lines if they prefer it.

from __future__ import annotations

from dataclasses import dataclass, field


@dataclass(eq=False)
class Node:
    val: int
    children: list[Node] = field(default_factory=list)


def compress(root: Node | None, compressed_nodes: list[int]) -> Node | None:
    """Remove every node whose value is in compressed_nodes (never the root).

    A removed node's children take its place in its parent's child list,
    in order. Changes the tree in place and returns the root.
    """
    if root is None:
        return None
    targets = set(compressed_nodes)
    targets.discard(root.val)  # the root is never compressed

    # replacement[node] = nodes that stand in node's place in its parent's list
    replacement: dict[int, list[Node]] = {}
    stack: list[tuple[Node, bool]] = [(root, False)]
    while stack:
        node, children_done = stack.pop()
        if not children_done:
            stack.append((node, True))
            for child in reversed(node.children):
                stack.append((child, False))
            continue
        new_children: list[Node] = []
        for child in node.children:
            new_children.extend(replacement.pop(id(child)))
        if node.val in targets:
            replacement[id(node)] = new_children
        else:
            node.children = new_children
            replacement[id(node)] = [node]
    return root


def height_after_compress(root: Node | None, compressed_nodes: list[int]) -> int:
    """Follow-up: height (edges on the longest root-to-leaf path) of the
    compressed tree, without changing or rebuilding it."""
    if root is None:
        return -1
    targets = set(compressed_nodes) - {root.val}
    best = 0
    stack: list[tuple[Node, int]] = [(root, 0)]
    while stack:
        node, depth = stack.pop()
        best = max(best, depth)
        for child in node.children:
            stack.append((child, depth if child.val in targets else depth + 1))
    return best


def to_nested(node: Node) -> tuple:
    """(val, [children...]) form, handy for asserts."""
    return (node.val, [to_nested(c) for c in node.children])


def build(spec: tuple) -> Node:
    val, kids = spec
    return Node(val, [build(k) for k in kids])

height_after_compress counts depth only on kept nodes: a compressed child sits at its parent's depth, so its children land one level below the parent.

• compress: O(n + k) time for n nodes and k listed values. Each node is pushed twice and appended to exactly one child list. O(n) extra space for the stack and the replacement map. The tree is changed in place.
• height_after_compress: O(n + k) time, O(n) space for the stack, no tree changes.
• Brute force for comparison: O(n * k) time.

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

• compress(T, [2]) gives 1 -> [5, 6, 3, 4], 3 -> [7], 7 -> [8, 9]. The promoted children take 2's position.
• compress(T, [3, 7]) gives 1 -> [2, 8, 9, 4], 2 -> [5, 6]. A chain of compressed nodes collapses two levels.
• compress(T, [5, 6]) gives 1 -> [2, 3, 4], 2 -> [], and the rest unchanged. Compressed leaves simply vanish.
• compress(T, [1]) leaves the tree unchanged. The root is never compressed.
• compress(T, []) and compress(T, [42]) leave the tree unchanged. Unknown values are ignored.
• compress(T, [2, 3, 4, 7]) gives 1 -> [5, 6, 8, 9]. Everything except the root and leaves is removed.
• compress(None, [1]) returns None.
• A path 1 -> 2 -> ... -> 100000 with [2..99999] compressed gives 1 -> [100000] without hitting the recursion limit.
• height_after_compress(T, []) is 3. height_after_compress(T, [3, 7]) is 2. height_after_compress(T, [2, 3, 7]) is 1.

1. Return only the height (or another statistic) without rebuilding the tree. Do one DFS that carries depth, and do not add 1 when stepping into a compressed node. This is height_after_compress above. The same pattern gives leaf count, width per level, or the new parent of each node.
2. Remove the whole subtree instead of promoting children. In the post-order step, a compressed node returns an empty list and its subtree is skipped entirely. A pre-order DFS that does not descend into targets is enough, and it is cheaper.
3. The root can be compressed too. The result is now a forest: return replacement[root], which is a list of roots. Ask whether the caller wants a forest or a single virtual root.
4. Compress requests arrive one at a time on a large tree. Store a parent pointer and the child's index (or use a linked list of children) so each removal splices in O(number of children) without a search. With an index map from value to node, there is no O(n) lookup per request. Mention that list splices in Python are O(len) and a doubly linked child list makes them O(1).
5. Choose the fewest nodes to compress so the height is at most H. This turns into a DP over the tree with a height budget per subtree. Talk through the state (node, depth it ends up at) and reconstructing the chosen set before coding.

Mid-level (L4)

• Asks where promoted children go and whether compressed nodes can be nested.
• Writes a correct recursive post-order solution that handles chains of compressed nodes and keeps the order.
• Tests the root, an empty list, a leaf, and a two-level chain.

Senior (L5)

• Gets to O(n) with a set lookup, and explains why post-order is needed (children must be resolved before the parent).
• Notices the recursion depth risk on a path-shaped tree and makes it iterative or explains the trade-off.
• Answers the height-only follow-up without building a new tree.

Staff (L6+)

• States the invariant (kept nodes keep their relative pre-order) and uses it to argue correctness.
• Discusses in-place changes versus returning a new tree, and what that means for callers holding node references.
• Handles the streaming and forest variants, and outlines the height-budget DP with its state and cost.