Bringing the grid back

After a blackout the control room closes one switch at a time, but every switch in a wave closes together. The question is not which substations come back, it is which come back together.

The problem

A regional grid is fed from one source substation. Each substation, once live, can energise the two written on its card — some cards name two, some one, some none. A switching wave takes a minute and every substation that can be energised in that minute is energised in it, so everything the source feeds comes back in wave two, everything those feed comes back in wave three, and so on.

The switching plan is stored the way the cards are written: an entry is [name, feed_a, feed_b], where feed_a and feed_b are the entries of the two substations this one energises, and None means no feeder. The whole plan is one such entry, or None if the region has no source.

Return the substation names grouped by wave, earliest first, and within a wave in the order they are reached across the grid.

Input. plan — the nested switching plan, or None.

Output. A list of waves, each a list of substation names.

Example.

plan = ["OTTERY",
          ["PENFOLD", ["WYCHAM", None, None], None],
          ["TARLING", None, ["MARSHEND", None, None]]]

  ->  [["OTTERY"], ["PENFOLD", "TARLING"], ["WYCHAM", "MARSHEND"]]

OTTERY energises PENFOLD and TARLING in wave two. PENFOLD then feeds only WYCHAM and TARLING only MARSHEND, so wave three holds two substations, not four.

A second example, a radial feed with no branching at all:

plan = ["ASHTON", None, ["BRIDLE", None, ["CULVER", None, None]]]
  ->  [["ASHTON"], ["BRIDLE"], ["CULVER"]]

Three substations, three waves, one name each. The number of waves is the length of the longest chain of cards, not the size of the grid.

Constraints.

• 0 <= substations <= 2000
• names are 1 to 8 uppercase letters and are unique
• the plan is None or a three-element entry

Hints

Approach

Brute force

Do one wave at a time: for wave d, walk the whole plan and collect every substation d cards from the source. Each wave is a full walk of 2000 entries and a radial grid needs 2000 waves, so that is up to 4 million entry visits to produce 2000 names.

The insight

A queue already holds one whole wave at a time — freeze its length before draining it, and the length you froze is that wave's width.

When the loop is about to start a wave, the queue holds exactly the substations energised in that wave and nothing else, because every push made while draining the previous wave went in behind them. So len(queue) at that instant is the wave's size. Read it lazily instead and the pushes keep the inner loop alive, the waves merge into one row, and the code still terminates and still returns something — which is what makes the bug hard to see.

Algorithm

1. Build the grid from the nested plan; an entry of None is no substation.
2. If there is no source, return an empty list of waves.
3. Put the source on a queue.
4. While the queue is not empty, record width = len(queue).
5. Pull width substations, appending each name to the current row and pushing the two it feeds, skipping the cards that are None.
6. Append the row to the answer and go round again.

Complexity

Time O(n) — each substation is pushed once and pulled once. Space O(w) where w is the widest wave; on a fully branching grid the last wave holds about half the substations, so 2000 of them peak near 1000 queued references.

Solution

"""Bringing the grid back — level order with the queue width frozen per wave."""

import sys
from collections import deque

sys.setrecursionlimit(20000)   # the plan nests as deep as the restoration is long


class Station:
    """One substation. `feed_a` and `feed_b` are the two it can energise."""

    def __init__(self, name):
        self.name = name
        self.feed_a = None
        self.feed_b = None


def wire(entry):
    """[name, feed_a, feed_b] nested lists, None for no feeder -> the source."""
    if entry is None:
        return None
    name, feed_a, feed_b = entry
    station = Station(name)
    station.feed_a = wire(feed_a)
    station.feed_b = wire(feed_b)
    return station


def solve(plan):
    source = wire(plan)
    if source is None:
        return []
    waves, queue = [], deque([source])
    while queue:
        # invariant: the queue holds exactly the stations energised this wave,
        # so its length now is the width of the wave about to be drained.
        width = len(queue)
        row = []
        for _ in range(width):
            station = queue.popleft()
            row.append(station.name)
            if station.feed_a is not None:
                queue.append(station.feed_a)
            if station.feed_b is not None:
                queue.append(station.feed_b)
        waves.append(row)
    return waves


NETWORK = ["OTTERY",
           ["PENFOLD", ["WYCHAM", None, None], None],
           ["TARLING", None, ["MARSHEND", None, None]]]

FULL = ["HALLOW",
        ["KEYNE", ["LANGTOFT", None, None], ["SWINDALE", None, None]],
        ["REDMIRE", ["SNAPE", None, None], ["TEALBY", None, None]]]

RADIAL = ["ASHTON", None, ["BRIDLE", None, ["CULVER", None, None]]]

TESTS = [
    ((NETWORK,), [["OTTERY"], ["PENFOLD", "TARLING"], ["WYCHAM", "MARSHEND"]]),
    ((FULL,), [["HALLOW"], ["KEYNE", "REDMIRE"],
               ["LANGTOFT", "SWINDALE", "SNAPE", "TEALBY"]]),
    ((RADIAL,), [["ASHTON"], ["BRIDLE"], ["CULVER"]]),
    ((["ISLAND", None, None],), [["ISLAND"]]),
    ((None,), []),
]

Pitfalls

• Draining the queue inside one row. Every name lands in the first wave and the answer is a single row wrapped in a list.
• Pushing None cards onto the queue. They come back off as substations, so the row gains a blank name or the next attribute lookup raises.
• Returning [[]] for a region with no source. No wave happened, so the answer is []; one empty row claims a wave took place with nothing in it.
• Assuming waves double. The example's third wave holds two substations after a wave of two, because those cards name one feeder each.

Variants

• Break-even on the market run — the same tree walked downwards route by route rather than across in waves.
• Traversal families — where level order sits against the three depth-first orders, and what it costs in memory.