2020/python/day22.py

125 lines
2.7 KiB
Python

import fileinput
import itertools
from pprint import pprint
from collections import deque
from typing import List, Deque, Dict, Tuple, Optional, Set
from copy import deepcopy
Deck = Deque[int]
Decks = Dict[int, Deck]
FrozenDeck = Tuple[Tuple[int, ...], Tuple[int, ...]]
def parse() -> Decks:
decks: Decks = {}
deck: Optional[Deck] = None
player: Optional[int] = None
for line in fileinput.input():
line = line.strip()
if line.startswith("Player "):
if deck and player:
decks[player] = deck
deck = deque()
player = int(line[len("Player ") : -1])
elif line.isdigit():
assert deck is not None
deck.append(int(line))
if deck and player:
decks[player] = deck
return decks
def part1_round(decks: Decks):
p1 = decks[1].popleft()
p2 = decks[2].popleft()
assert p1 != p2
if p1 > p2:
decks[1].extend((p1, p2))
else:
decks[2].extend((p2, p1))
def part1(decks: Decks) -> None:
i = 0
while decks[1] and decks[2]:
part1_round(decks)
i += 1
winner = 1 if decks[1] else 2
score = sum(card * (i + 1) for i, card in enumerate(reversed(decks[winner])))
print("Part 1:", score, "(player", winner, "won)")
def freeze_decks(decks: Decks) -> FrozenDeck:
return (
tuple(decks[1]),
tuple(decks[2]),
)
def part2_round(decks: Decks, history: Set[FrozenDeck]):
frozen = freeze_decks(decks)
if frozen in history:
# Player 1 wins if the round happened before
return 1
if not decks[1]:
return 2
if not decks[2]:
return 1
history.add(frozen)
p1 = decks[1].popleft()
p2 = decks[2].popleft()
assert(p1 != p2)
if len(decks[1]) >= p1 and len(decks[2]) >= p2:
new_decks = {
1: deque(list(decks[1])[:p1]),
2: deque(list(decks[2])[:p2]),
}
winner = part2_game(new_decks)
elif p1 > p2:
winner = 1
else:
winner = 2
assert(winner == 1 or winner == 2)
if winner == 1:
decks[1].extend((p1, p2))
elif winner == 2:
decks[2].extend((p2, p1))
def part2_game(decks: Decks) -> Optional[int]:
i = 0
history: Set[FrozenDeck] = set()
winner: Optional[int] = None
while winner is None:
winner = part2_round(decks, history)
i += 1
return winner
def part2(decks: Decks) -> None:
winner = part2_game(decks)
assert(winner is not None)
score = sum(card * (i + 1) for i, card in enumerate(reversed(decks[winner])))
print("Part 2:", score, "(player", winner, "won)")
def main() -> None:
decks = parse()
part1(deepcopy(decks))
part2(decks)
if __name__ == "__main__":
main()