refactor day 7 to shorten

1 files changed, 25 insertions, 37 deletions

diff --git a/7.py b/7.py
index 914cb30..58fdbd8 100644
--- a/7.py
+++ b/7.py
@@ -1,11 +1,11 @@
+# pyright: strict
 from collections import Counter
+from collections.abc import Callable
 from sys import stdin
 from typing import TypeAlias
 
 Hand: TypeAlias = tuple[int, int, int, int, int]
 
-inp = [(l[0], int(l[1])) for l in (l.rstrip().split() for l in stdin)]
-
 
 def parse_hand(card_worth: dict[str, int], hand: str) -> Hand:
     ret = tuple(card_worth[c] for c in hand)
@@ -30,48 +30,36 @@ def hand_type(c: list[tuple[int, int]]) -> int:
     return 0
 
 
-def p1(inp: list[tuple[str, int]]):
-    cards = "AKQJT98765432"
+def solve(
+    cards: str,
+    inp: list[tuple[str, int]],
+    fudge_count: Callable[[Counter[int]], Counter[int]] = lambda x: x,
+) -> int:
     card_worth = {c: i for i, c in enumerate(reversed(cards))}
 
-    def _hand_type(hand: Hand) -> int:
-        return hand_type(Counter(hand).most_common())
-
     return sum(
-        (i + 1) * h[1]
-        for i, h in enumerate(
+        rank * bid
+        for rank, (_, bid) in enumerate(
             sorted(
-                ((parse_hand(card_worth, i[0]), i[1]) for i in inp),
-                key=lambda i: (_hand_type(i[0]), *i[0]),
-            )
+                ((parse_hand(card_worth, hand), bid) for hand, bid in inp),
+                key=lambda i: (
+                    hand_type(fudge_count(Counter(i[0])).most_common()),
+                    i[0],
+                ),
+            ),
+            start=1,
         )
     )
 
 
-def p2(inp: list[tuple[str, int]]):
-    cards = "AKQT98765432J"
-    card_worth = {c: i for i, c in enumerate(reversed(cards))}
+def p2_fudge(c: Counter[int]) -> Counter[int]:
+    if 0 in c and c[0] != 5:
+        target = next(e[0] for e in c.most_common() if e[0] != 0)
+        c[target] += c[0]
+        c[0] = 0
+    return c
 
-    def _hand_type(hand: Hand) -> int:
-        c = Counter(hand)
-        mc = c.most_common()
-        if 0 in c and len(mc) != 1:
-            target = next(e[0] for e in mc if e[0] != 0)
-            c[target] += c[0]
-            c[0] = 0
-            mc = c.most_common()
-        return hand_type(mc)
 
-    return sum(
-        (i + 1) * h[1]
-        for i, h in enumerate(
-            sorted(
-                ((parse_hand(card_worth, i[0]), i[1]) for i in inp),
-                key=lambda i: (_hand_type(i[0]), *i[0]),
-            ),
-        )
-    )
-
-
-print(p1(inp))
-print(p2(inp))
+inp = [(l[0], int(l[1])) for l in (l.rstrip().split() for l in stdin)]
+print(solve("AKQJT98765432", inp))
+print(solve("AKQT98765432J", inp, fudge_count=p2_fudge))