Advent of Code 2019 (Day 6, Part 1)

AoC2019-06-1.py

"""
Advent of Code 2019: Day 6, Part 1
"""

from collections import defaultdict

tree = defaultdict(list)

with open("06-input.txt") as f:
    for line in f.readlines():
        tree[line[:3]].append(line[4:7])
 
# checksum = sum(k * #(nodes at depth k): nodes at depth k, k = 1, 2, ...)
# "Depth" means distance from root node ("COM").
checksum, depth = 0, 1
nodes = ["COM"]

# Prune nodes of increasing depth from tree.
while tree:
    next_nodes = []
    for node in nodes:
        children = tree.pop(node, [])
        checksum += depth * len(children)
        next_nodes.extend(children)
    nodes = next_nodes
    depth += 1
    
print(checksum)