Remove and Compact a Parent-Index Forest
Company: Pinterest
Role: Software Engineer
Category: Coding & Algorithms
Difficulty: medium
Interview Round: Technical Screen
## Problem
A forest is represented by `parent`, where node value is its array index. A root points to itself; every other node points to its parent's index. Remove a target node and all descendants, compact the surviving nodes in original index order, and rewrite parent indices to the compacted indices.
### Function Contract
Implement `remove_and_compact(parent, target) -> list[int]`. The input must not be mutated.
### Constraints
- `1 <= len(parent) <= 200000` and `0 <= target < len(parent)`.
- Every parent index is valid and the forest contains no cycle except root self-loops.
- Surviving roots must point to their own new indices.
- If no node survives, return an empty list.
### Examples
- `parent = [0,0,0,2,4,4]`, target `2` returns `[0,0,2,2]`.
- For `parent = [0,0,2]`, target `0`, return `[0]`; old node 2 becomes new root 0.
```hint Traverse downward
Construct child adjacency from the upward parent pointers, then mark the target's subtree.
```
```hint Assign every new index before rewriting
The surviving parent's new index must be known even when it appears later in a separate tree.
```
### Edge Cases
- The target can be a root, internal node, or leaf.
- Other trees in the forest keep their relative node order.
- Deleting one tree can change every later surviving index.
Overview: Remove a node and all descendants from a parent-index forest, preserve surviving order, compact indices, and rewrite every surviving parent reference without mutating the input.
A forest is represented by an integer array parent, where each node is identified by its array index. A root points to itself, and every other node points to its parent's index. Remove target and all of its descendants, then return the surviving parent array compacted in original index order with every parent index rewritten to the corresponding compacted index. Do not mutate the input. If no node survives, return an empty list.
Constraints
- 1 <= len(parent) <= 200000.
- 0 <= target < len(parent).
- Every parent index is valid, roots point to themselves, and the forest has no other cycle.
- Survivors retain original index order, and the input must not be mutated.
Examples
Input: ([0, 0, 0, 2, 4, 4], 2)
Expected Output: [0, 0, 2, 2]
Explanation: Removing node 2 deletes node 3 and remaps the surviving second tree.
Input: ([0, 0, 2], 0)
Expected Output: [0]
Explanation: Deleting the first tree leaves old node 2 as compacted root 0.
Hints
- Construct child adjacency from the upward parent pointers before traversing target's descendants.
- Assign every survivor's new index before rewriting parent pointers.
Community answers
Answer by richardxue0328
`
from collections import deque
def remove_and_compact(parent: list[int], target: int) -> list[int]:
n = len(parent)
# children[i] holds the direct children of i (root self-loops excluded).
children = [[] for _ in range(n)]
for node, par in enumerate(parent):
if par != node:
children[par].append(node)
# Iterative downward traversal marks the whole subtree at target.
removed = bytearray(n)
removed[target] = 1
queue = deque([target])
while queue:
node = queue.popleft()
for child in children[node]:
if not removed[child]:
removed[child] = 1
queue.append(child)
# Pass 1: assign every survivor its new index, in original order.
new_index = [-1] * n
survivors = []
for node in range(n):
if not removed[node]:
new_index[node] = len(survivors)
survivors.append(node)
# Pass 2: rewrite pointers. A survivor's parent always survives too,
# so every lookup is defined; roots map to their own new index.
return [new_index[parent[node]] for node in survivors]
`