Limit tree height by deleting nodes

The tree is given as children: dict[int, list[int]] and a root id. Height counts nodes on the longest root-to-leaf path. Deleting a node splices it out: its children attach to its nearest kept ancestor, so the descendants stay and only the deleted node stops counting.

Part 1: height after a given set of deletions.

You don't need to rebuild the tree. After the splice, each node's depth is the number of kept nodes on its original root path, including itself. Walk down from the root and carry a counter:

• kept_depth(child) = kept_depth(parent) + (0 if child is deleted else 1)
• height = max(kept_depth(v)) over all nodes.

A deleted leaf has the same kept depth as its parent, so it never raises the maximum. That matches the leaf disappearing. One pass is O(n). Use an explicit stack, because an interviewer may give a path-shaped tree with 10^5 nodes, and Python's recursion limit would fail on it.

Part 2: the fewest deletions that bring the height to at most k. This is the "limit" in the title, and it is the usual second part.

• Brute force: try every subset of non-root nodes, smallest first, and check each one with Part 1. That is O(2^n · n). It is only useful as a test oracle.

• Greedy pitfalls: "delete the deepest node on the longest path" and "delete the highest node on the longest path" are both easy to break. A high node fixes many paths at once, but it may sit above short branches that never needed fixing.

• Optimal: tree DP over a budget. Let dp[v][t] be the fewest deletions inside subtree(v) such that every downward path that starts at v has at most t kept nodes. At each node there are two choices:

  • keep v (needs t >= 1): sum(dp[c][t-1] for c in children(v))
  • delete v: 1 + sum(dp[c][t] for c in children(v))

  dp[v][t] is the smaller of the two. The root must be kept, so the answer is sum(dp[c][k-1] for c in children(root)). When k is at least the current height, the answer is 0.

Why the DP is correct. The paths that go through different children share no nodes below v, so their costs add up independently. Whether v is kept or deleted only changes the budget left for its children (t-1 or t). That is the exact sub-structure, and nothing is lost by solving each child on its own.

Invariant: dp[v][t] never increases as t grows. A larger budget never forces more deletions. So budgets above height(v) can be capped, which keeps the table at O(n · min(k, H)).

from typing import Dict, List, Set


def _postorder(children: Dict[int, List[int]], root: int) -> List[int]:
    """Nodes reachable from root, every child before its parent."""
    order, stack = [], [root]
    while stack:
        v = stack.pop()
        order.append(v)
        stack.extend(children.get(v, []))
    return order[::-1]


def height_after_deletions(children: Dict[int, List[int]], root: int,
                           deleted: Set[int]) -> int:
    """Part 1: height in nodes once every node in `deleted` is spliced out."""
    if root in deleted:
        raise ValueError("the root is never deleted")
    best = 0
    stack = [(root, 1)]
    while stack:
        v, depth = stack.pop()
        best = max(best, depth)
        for c in children.get(v, []):
            stack.append((c, depth + (0 if c in deleted else 1)))
    return best


def min_deletions_for_height(children: Dict[int, List[int]], root: int,
                             k: int) -> int:
    """Part 2: fewest non-root deletions that make the height <= k."""
    if k < 1:
        raise ValueError("the root alone already has height 1")
    order = _postorder(children, root)
    height: Dict[int, int] = {}
    for v in order:
        height[v] = 1 + max((height[c] for c in children.get(v, [])), default=0)
    if height[root] <= k:
        return 0

    # dp[v][t] for t in 0..cap(v); budgets above height(v) cost nothing extra.
    dp: Dict[int, List[int]] = {}
    for v in order[:-1]:        # every node except the root (last in postorder)
        kids = children.get(v, [])
        cap = min(k, height[v])
        row = [0] * (cap + 1)
        for t in range(cap + 1):
            delete = 1 + sum(_at(dp[c], t) for c in kids)
            keep = sum(_at(dp[c], t - 1) for c in kids) if t >= 1 else delete
            row[t] = min(keep, delete)
        for c in kids:          # free child rows once the parent is done
            del dp[c]
        dp[v] = row
    # The root is kept, so each child subtree gets a budget of k - 1.
    return sum(_at(dp[c], k - 1) for c in children.get(root, []))


def _at(row: List[int], t: int) -> int:
    """dp value at budget t; budgets past the row's cap cost the same as the cap."""
    return row[min(t, len(row) - 1)]

• height_after_deletions: O(n) time. O(n) space in the worst case for the stack, plus the deleted set. Lookups in the set are O(1).
• min_deletions_for_height: computing heights is O(n). The DP touches every parent-child edge once for each budget value, so it is O(n · min(k, H)) time, where H is the original height. Child rows are freed as soon as their parent is done, so live memory is bounded by the rows on the current frontier. The worst case is O(n · min(k, H)).
• Both functions are iterative, so a 10^5-node chain does not hit the recursion limit.

Tree T: 1 -> [2, 3, 4], 2 -> [5, 6], 3 -> [7], 7 -> [8] (height 4, via 1-3-7-8).

• height_after_deletions(T, 1, {}) → 4
• height_after_deletions(T, 1, {7}) → 3 (8 moves under 3)
• height_after_deletions(T, 1, {3, 7}) → 3 (path 1-2-5 is now the longest)
• height_after_deletions(T, 1, {2, 3, 7}) → 2 (everything hangs off the root)
• Single node {1: []}, no deletions → 1
• Deleting a leaf ({5}) does not change the height when another path is as long.
• min_deletions_for_height(T, 1, 4) → 0 (already fits)
• min_deletions_for_height(T, 1, 3) → 1 (delete 3 or 7)
• min_deletions_for_height(T, 1, 2) → 3 (delete 2 to fix both 5 and 6, then two of 3, 7, 8)
• min_deletions_for_height(T, 1, 1) → 7 (everything except the root)
• Chain of 100,000 nodes, k = 10 → 99990, with no recursion error.
• Randomized check: on 400 random trees with up to 10 nodes, the DP matches the brute force over subsets for every k.

1. Return which nodes to delete, not just how many. Store the choice (keep or delete) for each (v, t). Then walk down from the root with budget k - 1, passing t - 1 to a child's subtree when the child is kept and t when it is deleted. Child rows must be kept instead of freed, which costs O(n · min(k, H)) memory.
2. Deletions arrive as a stream, with a height query after each one. If there are few queries, rerunning Part 1 is O(n) each. For many queries, flatten the tree with an Euler tour. A deletion lowers the kept depth of the node's whole subtree by 1, so it becomes a range add of -1 on a contiguous interval. The height is a range max over the whole array. A lazy segment tree handles both in O(log n). Note that a deleted node still sits in the array at its parent's depth, which can't exceed the true maximum, so the answer stays correct.
3. Weighted nodes or deletion costs. If deleting node v costs w(v), replace 1 + with w(v) + in the delete branch. If height is a sum of node weights instead of a count, the budget dimension becomes a weight. The DP is then pseudo-polynomial, or you switch to a different formulation.
4. The input is a parent array or a list of edges. Build the children map in O(n) first. Find the root as the only node without a parent, and reject inputs with cycles or more than one root.
5. Deleting a node removes its whole subtree instead (no promotion). Part 1 becomes a DFS that skips deleted subtrees. Part 2 becomes trivial: every node deeper than k must go. The fewest delete operations is the number of nodes at depth k + 1, and the number of nodes removed is the count of nodes deeper than k. Ask which rule the interviewer means before you code.

Mid-level (L4)

• Asks whether height counts nodes or edges, whether the root can be deleted, and what happens to a deleted node's children.
• Solves Part 1 in O(n) by carrying a kept-depth counter, without rebuilding the tree.
• Writes the code cleanly and checks it on a small example plus the single-node case.

Senior (L5)

• Notices the recursion-depth risk in Python and writes an iterative traversal.
• Explains why simple greedy choices fail for Part 2, gives a counterexample, and reaches the dp[v][t] keep-or-delete recurrence.
• States the complexity precisely and caps budgets at the subtree height.

Staff (L6+)

• Justifies the DP with an argument about independent subtrees, and checks it against a brute force on random small trees.
• Handles the streaming follow-up with an Euler tour and a lazy segment tree, and explains why deleted nodes don't distort the max.
• Discusses input validation (cycles, missing ids, multiple roots) and how the API would change for a reconstruction of the deleted set or for weighted costs.