42 lines
1.1 KiB
Python
42 lines
1.1 KiB
Python
|
import itertools
|
||
|
import sys
|
||
|
from collections import defaultdict
|
||
|
from typing import Literal, cast
|
||
|
|
||
|
import parse # type: ignore
|
||
|
|
||
|
|
||
|
def max_change_in_happiness(happiness: dict[str, dict[str, int]]) -> int:
|
||
|
guests = list(happiness)
|
||
|
return max(
|
||
|
sum(
|
||
|
happiness[o][d] + happiness[d][o]
|
||
|
for o, d in zip((guests[0],) + order, order + (guests[0],))
|
||
|
)
|
||
|
for order in map(tuple, itertools.permutations(guests[1:]))
|
||
|
)
|
||
|
|
||
|
|
||
|
lines = sys.stdin.read().splitlines()
|
||
|
|
||
|
happiness: dict[str, dict[str, int]] = defaultdict(dict)
|
||
|
for line in lines:
|
||
|
u1, gain_or_loose, hap, u2 = cast(
|
||
|
tuple[str, Literal["gain", "lose"], int, str],
|
||
|
parse.parse( # type: ignore
|
||
|
"{} would {} {:d} happiness units by sitting next to {}.", line
|
||
|
),
|
||
|
)
|
||
|
happiness[u1][u2] = hap if gain_or_loose == "gain" else -hap
|
||
|
|
||
|
|
||
|
answer_1 = max_change_in_happiness(happiness)
|
||
|
print(f"answer 1 is {answer_1}")
|
||
|
|
||
|
for guest in list(happiness):
|
||
|
happiness["me"][guest] = 0
|
||
|
happiness[guest]["me"] = 0
|
||
|
|
||
|
answer_2 = max_change_in_happiness(happiness)
|
||
|
print(f"answer 2 is {answer_2}")
|