Treating a tree as a graph

"""Rooms at distance k, and the time for news to reach every room, by brute force; the rooms visited are counted."""
from collections import deque


class Node:
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None


def build(values):
    """The tree of a level-order list, where None marks a missing child."""
    if not values or values[0] is None:
        return None
    root = Node(values[0])
    waiting = deque([root])
    i = 1
    while i < len(values):
        node = waiting.popleft()
        for side in ("left", "right"):
            if i < len(values) and values[i] is not None:
                child = Node(values[i])
                setattr(node, side, child)
                waiting.append(child)
            i += 1
    return root


visits = 0


def path_to(root, room):
    """The rooms from the entrance down to room, found by searching the building from the top."""
    global visits
    above, stack = {root: None}, [root]
    while stack:
        node = stack.pop()
        visits += 1
        if node.val == room:
            path = []
            while node:
                path.append(node.val)
                node = above[node]
            return path[::-1]
        for child in (node.right, node.left):
            if child:
                above[child] = node
                stack.append(child)
    return None


def every_room(root):
    out, stack = [], [root]
    while stack:
        node = stack.pop()
        out.append(node.val)
        stack += [c for c in (node.right, node.left) if c]
    return out


def at_distance_brute(root, target, k):
    """For every room, find both paths from the entrance and measure the distance where they part."""
    to_target = path_to(root, target)
    found = []
    for room in every_room(root):
        to_room = path_to(root, room)
        shared = 0
        while shared < min(len(to_room), len(to_target)) and to_room[shared] == to_target[shared]:
            shared += 1
        if (len(to_room) - shared) + (len(to_target) - shared) == k:
            found.append(room)
    return sorted(found)


def minutes_brute(root, start):
    """Each minute, walk the whole building once and pass the news across every corridor next to someone who knows."""
    global visits
    heard, minutes, total = {start}, 0, len(every_room(root))
    while len(heard) < total:
        newly = set()
        stack = [(root, None)]
        while stack:
            node, above = stack.pop()
            visits += 1
            near = [c.val for c in (node.left, node.right) if c] + ([above] if above is not None else [])
            if node.val not in heard and any(r in heard for r in near):
                newly.add(node.val)
            stack += [(c, node.val) for c in (node.right, node.left) if c]
        heard |= newly
        minutes += 1
    return minutes


building = build([20, 11, 35, 4, 16, 28, 42, None, None, 13, 18, None, 31, None, None, None, None, 17])
for target, k in [(16, 2), (31, 3), (16, 0)]:
    visits = 0
    print(f"rooms {k} corridors from {target}: {at_distance_brute(building, target, k)}, {visits} rooms visited")
for start in (17, 20):
    visits = 0
    print(f"minutes for news from {start} to reach every room: {minutes_brute(building, start)}, {visits} rooms visited")

line = node = Node(0)
for v in range(1, 1000):
    node.left = Node(v)
    node = node.left
visits = 0
print(f"a line of 1,000, news from the far end: {minutes_brute(line, 999)} minutes, {visits:,} rooms visited")
visits = 0
at_distance_brute(line, 999, 10)
print(f"a line of 1,000, rooms 10 corridors from the far end: {visits:,} rooms visited")

rooms 2 corridors from 16: [4, 17, 20], 70 rooms visited
rooms 3 corridors from 31: [20, 42], 76 rooms visited
rooms 0 corridors from 16: [16], 70 rooms visited
minutes for news from 17 to reach every room: 7, 77 rooms visited
minutes for news from 20 to reach every room: 4, 44 rooms visited
a line of 1,000, news from the far end: 999 minutes, 999,000 rooms visited
a line of 1,000, rooms 10 corridors from the far end: 501,500 rooms visited