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
|
"""Advent of Code 2015, day 15"""
from collections import defaultdict
from itertools import combinations
from functools import cache
from typing import TypeVar, Generic
T = TypeVar('T')
class WeightedGraph(Generic[T]):
nodes: set[T]
weights: defaultdict[tuple[T, T], int]
def __init__(self):
self.nodes = set()
self.weights = defaultdict(int)
def add_edge(self, a, b, weight):
self.nodes.add(a)
self.nodes.add(b)
self.weights[tuple(sorted((a, b)))] = weight
def get_edge(self, a, b):
return self.weights[tuple(sorted((a, b)))]
def part1(graph: WeightedGraph[str]) -> int:
@cache
def max_happiness(first_guest: str, guest: str, remaining: frozenset[str]) -> int:
remaining = remaining - {guest}
if len(remaining) == 0:
return graph.get_edge(first_guest, guest)
return max(graph.get_edge(guest, other) + max_happiness(first_guest, other, remaining) for other in remaining)
return max(max_happiness(guest, guest, frozenset(graph.nodes)) for guest in graph.nodes)
def part2(graph: WeightedGraph[str]) -> int:
for guest in graph.nodes.copy():
graph.add_edge(guest, 'Me', 0)
return part1(graph)
if __name__ == '__main__':
guests = set()
edges = dict()
with open('input') as f:
for l in f:
a, _, what, how_much, _, _, _, _, _, _, b = l.rstrip()[:-1].split()
guests.add(a)
guests.add(b)
if what == 'lose':
how_much = -int(how_much)
edges[(a, b)] = int(how_much)
g = WeightedGraph()
for a, b in combinations(guests, 2):
g.add_edge(a, b, edges[(a, b)] + edges[(b, a)])
print(part1(g))
print(part2(g))
|
