aoc

ref: f09726471c0113b3a9e94d60e4537ef0fe2fca4f

006/main.py


 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
from functools import reduce


data = """COM)B
B)C
C)D
D)E
E)F
B)G
G)H
D)I
E)J
J)K
K)L"""


def add(a, b):
    return a + b


def assert_root(lines):
    orbitors = []
    orbitees = []

    for line in lines:
        a, b = line.split(')')
        orbitors.append(b)
        orbitees.append(a)

    orbitors = set(orbitors)
    orbitees = set(orbitees)

    root = orbitees.difference(orbitors)
    assert len(root) == 1
    assert 'COM' in root


def count(x, k, n=0):
    c = 0
    c += n

    for item in x.get(k, []):
        c += count(x, item, n + 1)

    return c


def main():
    x = {}

    data = open('input').read()

    lines = data.splitlines()

    assert_root(lines)

    for line in lines:
        a, b = line.split(')')

        if a not in x:
            x[a] = [b]
        else:
            xa = x[a]
            xa.append(b)
            x[a] = xa

    res = count(x, 'COM')
    print(res)



if __name__ == '__main__':
    main()