day7: [python] part 2

diff --git a/day7/run.py b/day7/run.py
index 80596e0775c7e714d644a4d59e34051827d10a12..89c7b8ffc105b2f57a3c5f30a776eca5faa0c61b 100755 (executable)
--- a/day7/run.py
+++ b/day7/run.py
@@ -3,9 +3,26 @@
 import sys
 import unittest
 import argparse
-from typing import Sequence
 from enum import IntEnum
 
+
+CARDVALUES = {
+    '2': '2',
+    '3': '3',
+    '4': '4',
+    '5': '5',
+    '6': '6',
+    '7': '7',
+    '8': '8',
+    '9': '9',
+    'T': 'A',
+    'J': 'B',
+    'Q': 'C',
+    'K': 'D',
+    'A': 'E'
+}
+
+
 class Win(IntEnum):
     HighCard = 1
     Pair = 2
@@ -15,32 +32,21 @@ class Win(IntEnum):
     FourKind = 6
     FiveKind = 7
 
-CARDVALUES = {
-    '2': '1',
-    '3': '2',
-    '4': '3',
-    '5': '4',
-    '6': '5',
-    '7': '6',
-    '8': '7',
-    '9': '8',
-    'T': '9',
-    'J': 'A',
-    'Q': 'B',
-    'K': 'C',
-    'A': 'D'
-}
 
-def hand_value(hand: str) -> str:
-    return "".join([CARDVALUES[card] for card in hand])
+def jokers(cards) -> int:
+    """Return the number of jokers"""
+    joker = [card for card in cards if card[0] == 'J']
+    if joker:
+        return joker[0][1]
+    return 0
+
 
-# 5, 4, fh, 3, 2p, p, flush
-def score_hand(s: str) -> Win:
+def score_hand(hand: str, part2=False) -> Win:
     hold = {
         '2': 0, '3': 0, '4': 0, '5': 0, '6': 0, '7': 0, '8': 0,
         '9': 0, 'T': 0, 'J': 0, 'Q': 0, 'K': 0, 'A': 0
     }
-    for c in s:
+    for c in hand:
         hold[c] += 1
     cards = [item for item in hold.items() if item[1] > 0]
     cards = sorted(cards, key=lambda x: x[1], reverse=True)
@@ -48,20 +54,50 @@ def score_hand(s: str) -> Win:
         win = Win.FiveKind
     elif cards[0][1] == 4:
         win = Win.FourKind
+        if part2 and jokers(cards):
+            win = Win.FiveKind
     elif cards[0][1] == 3:
         if cards[1][1] == 2:
             win = Win.FullHouse
+            if part2 and jokers(cards):
+                win = Win.FiveKind
         else:
             win = Win.ThreeKind
+            if part2:
+                j = jokers(cards)
+                if j == 3:
+                    win = Win.FourKind
+                elif j == 2:
+                    win = Win.FiveKind
+                elif j == 1:
+                    win = Win.FourKind
     elif cards[0][1] == 2:
         if cards[1][1] == 2:
             win = Win.TwoPair
+            if part2:
+                j = jokers(cards)
+                if j == 2:
+                    win = Win.FourKind
+                elif j == 1:
+                    win = Win.FullHouse
         else:
             win = Win.Pair
+            if part2 and jokers(cards):
+                win = Win.ThreeKind
     else:
         win = Win.HighCard
-    value = int(f"{win.value}" + hand_value(s), 16)
-    return value
+        if part2 and jokers(cards):
+            win = Win.Pair
+
+    # generate a hex number from the card faces and the hand type
+    # return this as a sortable integer value for the hand.
+    # For part2, jokers are lowest value
+    if part2:
+        CARDVALUES['J'] = '0'
+    value_str = "".join([CARDVALUES[card] for card in hand])
+    value = int(f"{win.value}{value_str}", 16)
+    return value, "".join([c[0] for c in cards]), win
+
 
 def load(filename):
     with open(filename) as stream:
@@ -69,24 +105,40 @@ def load(filename):
             hand, bid = line.strip().split(" ", maxsplit=1)
             yield hand, bid
 
+
 def main(args=None):
     parser = argparse.ArgumentParser(description="advent of code 2023 day 6")
     parser.add_argument('filename')
-    parser.add_argument('-t', '--test', action='store_true')
+    parser.add_argument('--debug', action='store_true')
+    parser.add_argument('--verbose', action='store_true')
+    parser.add_argument('--test', action='store_true')
     options = parser.parse_args(args)
 
     if options.test:
         return unittest.main()
-    
+
     total = 0
-    data = [(info[0], info[1], score_hand(info[0])) for info in load(options.filename)]
+    data = [(info[0], info[1], *score_hand(info[0], part2=False))
+            for info in load(options.filename)]
     for rank, info in enumerate(sorted(data, key=lambda x: x[2]), 1):
-        hand, bid, score = info
-        print(rank, hand, bid)
+        hand, bid, score, srt, win = info
+        if options.debug:
+            print(rank, hand, bid, hex(score))
         total += rank * int(bid)
-
     print(f"part 1: {total}")
-    #print(f"part 2: {part2(options.filename)}")
+
+    total2 = 0
+    data2 = [(info[0], info[1], *score_hand(info[0], part2=True))
+             for info in load(options.filename)]
+    for rank, info in enumerate(sorted(data2, key=lambda x: x[2]), 1):
+        hand, bid, score, srt, win = info
+        if options.debug:
+            if options.verbose:
+                print(rank, hand, bid, srt, win, hex(score))
+            else:
+                print(rank, hand, bid)
+        total2 += rank * int(bid)
+    print(f"part 2: {total2}")
 
 
 if __name__ == '__main__':