124 lines
2.7 KiB
Python
124 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()
|