add solution to part b

4/scratchcards.py

@@ -12,17 +12,22 @@ def get_winning_numbers(input_data: str) -> list:
     """
 
     # separate cards into a list of strings and remove excess whitespace
-    pairs = [ " ".join(card.split()) for card in input_data.split("\n") ]
+    pairs = [" ".join(card.split()) for card in input_data.split("\n")]
     # extract a set of winning numbers
-    winners = [ set(pair.split(":")[1].split("|")[0].strip().split(" ")) for pair in pairs ]
+    winners = [
+        set(pair.split(":")[1].split("|")[0].strip().split(" ")) for pair in pairs
+    ]
     # extract a set of scratch numbers
-    numbers = [ set(pair.split(":")[1].split("|")[1].strip().split(" ")) for pair in pairs ]
+    numbers = [
+        set(pair.split(":")[1].split("|")[1].strip().split(" ")) for pair in pairs
+    ]
     # calculate set intersection between numbers and winning numbers
-    intersections = [ numbers[i].intersection(winners[i]) for i in range(len(numbers)) ]
+    intersections = [numbers[i].intersection(winners[i]) for i in range(len(numbers))]
 
     return intersections
 
-def calculate_score(winning_numbers:list) -> list:
+
+def calculate_scores(winning_numbers: list) -> list:
     """
     Calculate the score for each number
     - 0 points for no match
@@ -36,10 +41,41 @@ def calculate_score(winning_numbers:list) -> list:
         if len(win) == 1:
             scores.append(1)
         if len(win) > 1:
-            scores.append(2**(len(win) - 1))
+            scores.append(2 ** (len(win) - 1))
+
     return scores
 
 
+def new_rules(input_data: list) -> list:
+    """
+    Calculate the score for the updated rule set.
+    - If your card has n matches, get the next n cards
+    - repeat until there are no more matches
+    """
+    data = input_data.split("\n")
+    score, win, total = 0, 0, 0
+    cardcounts = []
+
+    for i in range(len(data)):
+        cardcounts.append(1)
+    for x, row in enumerate(data):
+        wins = row.split("|")[0].split(":")[1].split()
+        nums = row.split("|")[1].split()
+        for num in nums:
+            if num in wins:
+                win += 1
+                if score == 0:
+                    score = 1
+                else:
+                    score *= 2
+        for i in range(win):
+            cardcounts[x + i + 1] += cardcounts[x]
+        total += score
+        score, win = 0, 0
+
+    return cardcounts
+
+
 if __name__ == "__main__":
     # get puzzle and parse data
     puzzle = Puzzle(year=2023, day=4)
@@ -47,11 +83,11 @@ if __name__ == "__main__":
     # parse_input(puzzle.input_data)
     winning_numbers = get_winning_numbers(puzzle.input_data)
 
-    answer_a = sum(calculate_score(winning_numbers))
+    answer_a = sum(calculate_scores(winning_numbers))
     print(f"{answer_a}")
     submit(answer_a, part="a", day=4, year=2023)
 
     # part b:
-    answer_b = 0
+    answer_b = sum(new_rules(puzzle.input_data))
     print(f"{answer_b}")
-    #submit(answer_b, part="b", day=4, year=2023)
+    submit(answer_b, part="b", day=4, year=2023)