python(day22): Do day 22

python/day22.py

@@ -0,0 +1,124 @@
+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()