~/problems / Graphs / DFS and connected components

Cutting a branch out of a forest

easy ~15 min Pinterest

A forest of n nodes, numbered 0 to n - 1, is stored as a parent array: parent[i] is the parent of node i, and a root is marked by being its own parent (parent[i] == i). There may be several roots. Indexes don't follow any order: a node's parent can have a larger or smaller number than the node itself.

Write

def delete_subtree(parent: list[int], node: int) -> list[int]

that deletes node together with everything below it (its children, their children, and so on) and returns the new parent array. Deleted nodes keep their positions, but their entry becomes -1. Every other entry stays exactly as it was. Don't modify the input list.

  • 1 <= n <= 200,000 and 0 <= node < n. The input is a valid forest with no -1 entries.
  • Deleting a root removes its whole tree; the other trees are untouched.
#   tree A:   3            tree B:  5
#            / \                    |
#           0   4                   2
#               |
#               1
parent = [3, 4, 5, 3, 3, 5]
delete_subtree(parent, 4)   # [3, -1, 5, 3, -1, 5]
delete_subtree(parent, 3)   # [-1, -1, 5, -1, -1, 5]
delete_subtree(parent, 2)   # [3, 4, -1, 3, 3, 5]

The tests include a single chain of 200,000 nodes. Walking up from every node to check whether it sits under node is quadratic there, and recursing down the chain overflows Python's call stack.

Show hint

Build a children list from parent (skipping each root's link to itself), then walk down from node with an explicit stack, marking everything you reach.

Topic: DFS and connected components. Iterative DFS / flood fill; count and label components.

Read the visual guide
0:00
Ctrl ' run · Ctrl ↵ submit
esc