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
|
from itertools import permutations, islice
import collections
def sliding_window(iterable, n):
# sliding_window('ABCDEFG', 4) -> ABCD BCDE CDEF DEFG
it = iter(iterable)
window = collections.deque(islice(it, n), maxlen=n)
if len(window) == n:
yield tuple(window)
for x in it:
window.append(x)
yield tuple(window)
def part1(nodes: set, distances: dict) -> int:
"""
>>> part1(set('ABC'), {('A', 'B'): 464, ('A', 'C'): 518, ('B', 'C'): 141})
605
"""
return min(sum(distances[tuple(sorted(edge))] for edge in sliding_window(path, 2)) for path in permutations(nodes))
def part2(nodes: set, distances: dict) -> int:
"""
>>> part2(set('ABC'), {('A', 'B'): 464, ('A', 'C'): 518, ('B', 'C'): 141})
982
"""
return max(sum(distances[tuple(sorted(edge))] for edge in sliding_window(path, 2)) for path in permutations(nodes))
if __name__ == '__main__':
nodes = set()
distances = dict()
with open('input') as f:
for l in f:
l = l.strip()
edge, distance = l.split(' = ')
edge = tuple(sorted(edge.split(' to ')))
nodes.add(edge[0])
nodes.add(edge[1])
distances[edge] = int(distance)
print(part1(nodes, distances))
print(part2(nodes, distances))
|
