Peel an N-ary tree one leaf layer per day: Frontier Interview's solution

Question: Peel an N-ary tree one leaf layer per day

Companies: Google

Difficulty: Medium

Problem statement

Every day, all leaves of an N-ary tree fall off at the same time. A leaf is a node with no children left. When a node loses its last child, it becomes a leaf and falls on the next day, not the same day. This repeats until the tree is gone. The root always falls last.

Implement:

class Node:
    val: int
    children: list["Node"]

def leaves_by_day(root: Node | None) -> list[list[int]]: ...

Return days, where days[i] lists the values of the nodes removed on day i + 1. Values inside one day may be in any order. Return [] for an empty tree.

Constraints

0 <= n <= 10^5 nodes. The tree can be very deep (a single chain).
Values are integers and may repeat. Treat every node as distinct.
You do not need to change the tree. Just report the removal schedule.

Example 1

        10
      /  |  \
    20  30   40
   /  \      |
  50  60     70

Output: [[30, 50, 60, 70], [20, 40], [10]]. On day 1, 30, 50, 60 and 70 are leaves. On day 2, 20 and 40 have no children left. On day 3, only 10 is left.

Example 2

A chain 1 -> 2 -> 3 (each node has one child). Output: [[3], [2], [1]].

Reconstructed from the reported question; details may differ from what interviewers ask.

Approach

Key idea: the day a node is removed equals its height. Define the height
of a leaf as 0 and the height of any other node as 1 + max(height(child)).
A leaf goes on day 1 (index 0). A parent becomes a leaf only after its last
child is gone, and its last child is the one with the greatest height. So the
node is removed one day after its tallest child: day(node) = 1 + max(day(child)).
That is exactly the height recurrence. The root has the greatest height, so
it is always removed on the last day, alone.

Brute force (where many candidates start). Simulate the days. Each day,
walk the tree, collect every node with no remaining children, then detach
them from their parents. A path-shaped tree (a chain of n nodes) takes n
days and each day walks the remaining tree, so this is O(n^2). Say it, then
improve it.

Optimal: one post-order pass. Compute every node's height after its
children are done, and append the node's value to days[height]. Because a
node's height is computed only after all of its children, days grows by at
most one list at a time. I use an iterative post-order so a deep chain (say
10^5 nodes) does not overflow Python's recursion limit.

Alternative: peel the leaves (Kahn style). Count each node's remaining
children and keep a parent pointer. Start a queue with all leaves. Process the
queue one level at a time: every node in the current level is removed today;
decrement its parent's count and, when it reaches zero, the parent joins the
next level. This mirrors the problem statement literally, also runs in O(n),
and is the natural base for follow-ups that change the removal rule.

Invariants

height(node) > height(child) for every child, so each node lands on a
strictly later day than all of its children.
Every day from 0 to height(root) is non-empty: a node of height h > 0
has a child of height h - 1.
The tree is not modified. Removal is computed, not performed, unless the
interviewer asks for real deletion (see follow-ups).

Code

from __future__ import annotations

from collections import deque
from dataclasses import dataclass, field
from typing import Optional


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


def leaves_by_day(root: Optional[Node]) -> list[list[int]]:
    """days[i] holds the values removed on day i + 1 (height-based, O(n))."""
    if root is None:
        return []
    days: list[list[int]] = []
    height: dict[int, int] = {}          # id(node) -> height
    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 node.children:
                stack.append((child, False))
            continue
        h = 0
        if node.children:
            h = 1 + max(height[id(c)] for c in node.children)
        height[id(node)] = h
        if h == len(days):
            days.append([])
        days[h].append(node.val)
    return days


def leaves_by_day_peeling(root: Optional[Node]) -> list[list[int]]:
    """Same answer by simulating the days with child counts (Kahn style)."""
    if root is None:
        return []
    parent: dict[int, Optional[Node]] = {id(root): None}
    remaining: dict[int, int] = {}
    order = [root]
    for node in order:                    # BFS that grows while iterating
        remaining[id(node)] = len(node.children)
        for child in node.children:
            parent[id(child)] = node
            order.append(child)
    today = deque(n for n in order if not n.children)
    days: list[list[int]] = []
    while today:
        removed, tomorrow = [], deque()
        for node in today:
            removed.append(node.val)
            p = parent[id(node)]
            if p is not None:
                remaining[id(p)] -= 1
                if remaining[id(p)] == 0:
                    tomorrow.append(p)
        days.append(removed)
        today = tomorrow
    return days


def _norm(days: list[list[int]]) -> list[list[int]]:
    return [sorted(d) for d in days]


if __name__ == "__main__":
    #        1
    #     /  |  \
    #    2   3   4
    #   / \      |
    #  5   6     7
    #            |
    #            8
    tree = Node(1, [
        Node(2, [Node(5), Node(6)]),
        Node(3),
        Node(4, [Node(7, [Node(8)])]),
    ])
    expected = [[3, 5, 6, 8], [2, 7], [4], [1]]
    assert _norm(leaves_by_day(tree)) == expected
    assert _norm(leaves_by_day_peeling(tree)) == expected

    assert leaves_by_day(None) == [] and leaves_by_day_peeling(None) == []
    assert leaves_by_day(Node(9)) == [[9]]

    chain = Node(0)                       # deep chain: no recursion limit
    cur = chain
    for i in range(1, 100_000):
        nxt = Node(i)
        cur.children.append(nxt)
        cur = nxt
    days = leaves_by_day(chain)
    assert len(days) == 100_000 and days[0] == [99_999] and days[-1] == [0]

    star = Node(0, [Node(i) for i in range(1, 6)])
    assert _norm(leaves_by_day(star)) == [[1, 2, 3, 4, 5], [0]]
    print("all tests passed")

Complexity

leaves_by_day: O(n) time. Each node is pushed twice and each child edge
is read twice (once to push, once in max). O(n) space for the stack, the
height map and the output.
leaves_by_day_peeling: O(n) time. Each node is enqueued once and each
edge decrements one counter once. O(n) space for the parent map, counts,
BFS order and output.
Brute-force simulation: O(n * D) time where D is the number of days
(O(n^2) for a chain), O(n) space.

Test cases

Empty tree (root = None) -> [].
Single node 9 -> [[9]].
Star: root 0 with children 1..5 -> [[1, 2, 3, 4, 5], [0]].
Uneven tree from the code: root 1 with children 2 (children 5, 6),
3, and 4 (child 7, which has child 8) ->
[[3, 5, 6, 8], [2, 7], [4], [1]]. Node 2 goes on day 2 even though its
sibling 4 is still there; node 4 waits for the deeper branch.
Chain 0 -> 1 -> ... -> 99999 -> 100000 days, one node per day, deepest
first, root last. Checks there is no recursion overflow.
Duplicate values (two different nodes both valued 7) -> both appear, on
their own days. The code keys on node identity, not value.
Root is always alone in the last list, for any non-empty tree.

Follow-ups

Actually delete the nodes. Return the days and leave the tree empty, or
stop after k days and return the remaining tree. With heights known,
keep a node only if height > k - 1 and filter each child list in the same
post-order pass; set children of removed nodes to [] so they can be
freed.
The tree is given as a parent array or edge list, not node objects.
Build child counts from the parent array and run the peeling version with
integer ids. No recursion, no id(), and it handles 10^6 nodes easily.
Undirected tree with no root (peel all degree-1 nodes each day). The
peeling version still works with degree instead of child count. The last
one or two nodes left are the tree's centers (the "minimum height trees"
problem).
Each node has a lifetime: it falls t[node] days after it becomes a
leaf. Now day(node) = t[node] + max(day(child)) (with 0 for no
children); it is still one post-order pass. If several nodes can be removed
on fractional or arbitrary times, collect (time, val) and sort or bucket
by time.
Return only the day a given node is removed, or answer many such
queries. That is the node's height. Precompute all heights once in
O(n) and answer each query in O(1). If the tree changes (a leaf is
added), only heights on the path to the root can change; walk up while
the height increases.

Interviewer expectations

Mid-level (L4)

Gets a correct answer, even by simulating day by day first, and states
that the simulation is O(n^2) on a chain.
Moves to the height observation with a hint and writes a recursive
post-order version that produces the right groups.
Tests a leaf-only tree, a star, and an uneven tree.

Senior (L5)

Explains why removal day equals height (last child to go is the tallest)
without prompting, and gives O(n) time and space.
Notices that recursion can overflow on a deep chain and writes an iterative
version, or names the limit and the fix.
Handles duplicate values by keying on node identity, and covers the empty
tree.

Staff (L6+)

Offers both views (height and leaf peeling) and picks the one that fits the
next part: peeling for changed removal rules, unrooted trees or edge-list
input; height for per-node queries.
Handles the follow-ups quickly: real deletion after k days, weighted
lifetimes, incremental updates to heights.
Discusses scale: integer ids and arrays instead of objects and hash maps for
millions of nodes, and the memory cost of keeping the whole output.