summaryrefslogtreecommitdiffstats
path: root/20/py/d10.py
blob: 2b3bc2ebbded7adf5aa5a31c7e5dae69ba2c5446 (plain) (blame)
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
#!/usr/bin/env python3
import aoc20
import functools
import sys


def pt1(_in):
    jolts = [int(n.strip()) for n in _in]
    jolts.append(max(jolts)+3)
    jolts = sorted(jolts)
    jolt = 0
    diffs = [0, 0, 0]
    for j in jolts:
        diffs[j-jolt-1] += 1
        jolt = j
    return diffs[0] * diffs[2]


def pt2(_in):
    jolts = [int(n.strip()) for n in _in]
    jolts.append(max(jolts)+3)
    jolts = sorted(jolts)
    jolts.insert(0, 0)
    jolt = 0
    reach = {}
    for i in range(len(jolts)):
        reachable = []
        j = i+1
        while j < len(jolts) and jolts[j] - jolts[i] <= 3:
            reachable.append(jolts[j])
            j += 1
        reach[jolts[i]] = tuple(reachable)

    @functools.cache
    def ways(start, dest):
        if start == dest:
            return 1
        return sum(ways(child, dest) for child in reach[start])

    return ways(jolts[0], jolts[-1])


if __name__ == "__main__":
    input = aoc20.read_input(sys.argv[1:], 10)
    print(pt1(input))
    print(pt2(input))