~/problems / Trees / Tree DP

Count subordinates in a company tree

easy ~15 min

Implement subordinates(n, bosses) -> list[int].

A company has n employees numbered 1..n. Employee 1 is the director and has no boss. bosses has length n - 1: bosses[i] is the direct boss of employee i + 2. The reporting lines form a tree rooted at employee 1. A boss's number can be larger than their employee's number.

Return a list of length n whose k-th entry (0-based) is the number of subordinates of employee k + 1: everyone below them in the tree, direct or indirect.

subordinates(5, [1, 1, 2, 3])
# employee 2's boss is 1, 3's boss is 1, 4's boss is 2, 5's boss is 3
# [4, 1, 1, 0, 0]

n goes up to 100,000, and the tree can be a single long chain, so a recursive solution hits Python's recursion limit. Aim for O(n): walking up from every employee to the root is O(n²) on a chain.

Show hint

An employee's count is easy once all of their direct reports' counts are known. Find an order of the employees in which everyone comes after their boss (without recursion), then process it in reverse.

Topic: Tree DP. Post-order: each node returns a small tuple of states to its parent.

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