advent of code day 4 part 2

aoc_2023_day_4_part_2.py

from collections import defaultdict
from dataclasses import dataclass
from typing import List, Set


@dataclass
class Deck:
    id: int
    winning: Set[int]
    mine: Set[int]

    def count_matches(self):
        return len(self.winning.intersection(self.mine))


@dataclass
class DecksWithCopies:
    decks: List[Deck]


def parse_raw_cards(raw_cards: str) -> List[int]:
    return [
        (int(c)) for c
        in raw_cards.strip().split(' ')
        if c != ''
    ]


def solve(input_string):
    original_deck = []

    for line in input_string.splitlines():
        deck, cards = line.split(':')
        deck_id = deck.split(' ')[-1]
        raw_winning, raw_mine = cards.split('|')

        original_deck.append(Deck(
            id=int(deck_id),
            winning=set(parse_raw_cards(raw_winning)),
            mine=set(parse_raw_cards(raw_mine))
            )
        )

    deck_with_copies = []
    carry = defaultdict(int)  # key = deck_id, value = copies count
    for idx, deck in enumerate(original_deck):
        copy_for_current_deck = carry.pop(deck.id, 0)

        current_decks = DecksWithCopies(
            decks=[deck] + ([deck] * copy_for_current_deck)
        )

        current_total_matches = current_decks.decks[0].count_matches()
        for i in range(deck.id+1, deck.id + 1 + current_total_matches):
            carry[i] += len(current_decks.decks)

        deck_with_copies.append(current_decks)

    return sum(len(decks.decks)for decks in deck_with_copies)


if __name__ == '__main__':
    print(solve(_input))