The storm-drain survey

A borough surveys its storm drains before the winter. The camera is hired by the day, so the crew wants as few visits as possible.

The problem

The drains run from one outfall. Every chamber but that one drains into a chamber nearer it, so the pipes branch and never rejoin.

A camera lowered into a chamber sees that chamber and, along each pipe meeting it, the chamber at the far end — one pipe, no further. Every chamber must be seen; report the fewest visits.

Input. drains — a list of integers, where drains[i] is the chamber that chamber i drains into. Exactly one entry is -1, marking the outfall.

Output. The smallest number of chambers to lower the camera into.

Example.

drains = [-1, 0, 1, 1, 3, 3]   ->  2

Chamber 1 feeds the outfall and takes 2 and 3; chamber 3 takes 4 and 5. A camera in 1 sees 0, 1, 2 and 3; one in 3 sees 4 and 5. One camera cannot do it — chambers 2 and 4 lie three pipes apart.

A second example, where the obvious first choice is wrong:

drains = [-1, 0, 0, 1, 1, 2, 2]   ->  2

The outfall looks like the place to start: a camera there sees three chambers. But it leaves the four dead ends unseen and costs three visits, where cameras in 1 and 2 cover all seven.

A third example, a drain with no branching at all:

drains = [-1, 0, 1, 2, 3, 4, 5]   ->  3

Cameras in chambers 1 and 4 see the first six between them, and the last chamber needs a third. Taking every other chamber from the outfall costs four.

Constraints.

• 1 <= len(drains) <= 10^4
• exactly one entry is -1; every other entry is a valid index
• the pipes form one connected network with no loops

Hints

Approach

Brute force

Try every set of chambers and check each one is seen: 2ⁿ sets. Twenty-five chambers is 33 million; the network holds ten thousand.

The insight

One number cannot score a branch: the chamber above pays a different price depending on whether the branch covers itself and whether it covers that chamber — so carry all three prices up.

Write camera[c], watched[c] and blind[c] for the fewest visits in c's branch when the camera goes into c, when c is seen from below, and when everything under c is seen but c is not. A dead end starts at 1, unreachable and 0. Each is a sum over the chambers feeding c; only watched[c] carries a condition — one of them must hold a camera — so take the cheapest everywhere, then pay the smallest upgrade if none does.

Algorithm

1. Build the list of chambers feeding each chamber, and find the outfall.
2. Flatten the network with a stack, parents before children.
3. Walk it backwards, filling a chamber's three values once every chamber feeding it is settled.
4. Answer: the smaller of camera[outfall] and watched[outfall].

Complexity

Time O(n) — every chamber and pipe touched a constant number of times. Space O(n) for the three arrays, the feed lists and the order.

Solution

"""The storm-drain survey — three answers per chamber, settled from the ends inward."""

BIG = float('inf')


def solve(drains):
    n = len(drains)
    below = [[] for _ in range(n)]
    root = 0
    for chamber, downstream in enumerate(drains):
        if downstream < 0:
            root = chamber
        else:
            below[downstream].append(chamber)

    # Flatten first: a network can be one long run, and a recursive sweep
    # would hit the interpreter's depth limit before it hit a wrong answer.
    order, stack = [], [root]
    while stack:
        chamber = stack.pop()
        order.append(chamber)
        stack.extend(below[chamber])

    # camera[c]  — fewest cameras in c's branch with one lowered into c
    # watched[c] — fewest with c watched from a chamber above it in the branch
    # blind[c]   — fewest with everything under c watched but c itself not
    camera = [1] * n
    watched = [BIG] * n
    blind = [0] * n
    for chamber in reversed(order):       # children settle before their parent
        kids = below[chamber]
        if not kids:
            continue                      # a dead end keeps the leaf values
        camera[chamber] = 1 + sum(min(camera[k], watched[k], blind[k]) for k in kids)
        blind[chamber] = sum(watched[k] for k in kids)
        base = sum(min(camera[k], watched[k]) for k in kids)
        # to be watched without its own camera, at least one child needs one
        surplus = min(camera[k] - min(camera[k], watched[k]) for k in kids)
        watched[chamber] = base + surplus
    return min(camera[root], watched[root])

TESTS = [
    (([-1, 0, 1, 1, 3, 3],), 2),
    (([-1, 0, 0, 1, 1, 2, 2],), 2),
    (([-1, 0, 1, 2, 3, 4, 5],), 3),
    (([-1],), 1),
    (([-1, 0],), 1),
    (([1, -1],), 1),
    (([-1, 0, 0, 0, 0],), 1),
    (([-1, 0, 0, 1, 1, 2, 2, 3, 3, 4],), 3),
    (([-1] + list(range(0, 999)),), 334),
]

Pitfalls

• Scoring a branch with one number. In [-1, 0, 0, 1, 1, 2, 2] the branches under chambers 1 and 2 cost one camera each, and those cameras see the outfall too. A sweep that records only the cost cannot know that, and books a third.
• Letting an unseen chamber cost nothing. blind[c] is available only when the chamber above takes a camera, so the outfall may never end on it: read as a third option at the top, [-1] answers 0.
• Sweeping with recursion. A drain can run in one line for ten thousand chambers, and the stack gives out before the answer does.

Variants

• Cutting the parcel spur — one number per node is enough there: each subtree is asked only for a total.
• Tree, digit and bitmask DP — the lesson behind an answer assembled subtree by subtree.