Rerooting: answers for every root

free/brute_force.py

"""The total distance from every village to all others by brute force: one breadth-first search per village."""
import random

from villages import distances_from, neighbours, random_district


def totals_brute(n, roads):
    near = neighbours(n, roads)
    visits = 0
    totals = []
    for start in range(n):
        dist = distances_from(near, start)
        visits += n  # one search reaches every village once
        totals.append(sum(dist))
    return totals, visits


DISTRICTS = {
    "the seven villages": (7, [(0, 1), (1, 2), (1, 3), (3, 4), (3, 5), (5, 6)]),
    "a line of four villages": (4, [(0, 1), (1, 2), (2, 3)]),
    "a star of five villages around village 0": (5, [(0, 1), (0, 2), (0, 3), (0, 4)]),
    "one village": (1, []),
}
for label, (n, roads) in DISTRICTS.items():
    totals, visits = totals_brute(n, roads)
    print(f"{label}: {totals} after {visits:,} village visits")

rng = random.Random(14)
for n in (500, 1_000, 2_000):
    _, visits = totals_brute(n, random_district(rng, n))
    print(f"a random district of {n:,} villages: {visits:,} village visits")

the seven villages: [15, 10, 15, 9, 14, 12, 17] after 49 village visits
a line of four villages: [6, 4, 4, 6] after 16 village visits
a star of five villages around village 0: [4, 7, 7, 7, 7] after 25 village visits
one village: [0] after 1 village visits
a random district of 500 villages: 250,000 village visits
a random district of 1,000 villages: 1,000,000 village visits
a random district of 2,000 villages: 4,000,000 village visits

free/villages.py

"""Districts of villages joined by roads of 1 km that form a tree: neighbour lists, searches and random districts."""
from collections import deque


def neighbours(n, roads):
    """The villages one road away from each village."""
    near = [[] for _ in range(n)]
    for u, v in roads:
        near[u].append(v)
        near[v].append(u)
    return near


def distances_from(near, start):
    """Breadth-first search from one village: the distance in km to every village."""
    dist = [-1] * len(near)
    dist[start] = 0
    queue = deque([start])
    while queue:
        u = queue.popleft()
        for v in near[u]:
            if dist[v] < 0:
                dist[v] = dist[u] + 1
                queue.append(v)
    return dist


def random_district(rng, n):
    """n villages with random names 0..n-1; each joins an earlier one; roads in random order and direction."""
    label = list(range(n))
    rng.shuffle(label)
    roads = [(label[rng.randrange(v)], label[v]) for v in range(1, n)]
    roads = [(v, u) if rng.random() < 0.5 else (u, v) for u, v in roads]
    rng.shuffle(roads)
    return roads