Limit tree height by deleting nodes

You are given a rooted N-ary tree. Every node has a unique integer id. The tree is passed as children: dict[int, list[int]] (a node with no entry has no children) together with the root id.

Deleting a node splices it out. Its children, and their whole subtrees, attach to the deleted node's parent. If the parent was deleted too, they attach to the nearest ancestor that was kept. Descendants are never removed. The root is never deleted.

Height is the number of nodes on the longest root-to-leaf path. A tree with only the root has height 1.

Part 1

def height_after_deletions(children: dict[int, list[int]], root: int, deleted: set[int]) -> int

Return the height of the tree after every node in deleted is spliced out.

Part 2

def min_deletions_for_height(children: dict[int, list[int]], root: int, k: int) -> int

Return the smallest number of non-root nodes to delete so that the height becomes at most k (k >= 1).

Constraints

• 1 <= number of nodes <= 10^5. The tree may be a long chain, so avoid deep recursion.
• 1 <= k <= number of nodes.

Example

1 -> [2, 3, 4]
2 -> [5, 6]
3 -> [7]
7 -> [8]

• With no deletions, the height is 4 (path 1-3-7-8).
• deleted = {7}: node 8 moves under 3, so the height is 3.
• deleted = {2, 3, 7}: nodes 5, 6, 8 and 4 all hang directly off the root, so the height is 2.
• min_deletions_for_height(..., k=3) is 1 (delete 3 or 7).
• min_deletions_for_height(..., k=2) is 3. One valid answer is {2, 3, 7}.