advent-of-code-2023/day7/day7.py

import sys
sys.path.append('../lib')

import tools
from functools import total_ordering

HIGH = 1
PAIR = 2
TPAIR = 3
THREE = 4
FHOUSE = 5
FOUR = 6
FIVE = 7 


DEFAULT_DECK = ['2', '3', '4', '5', '6', '7', '8', '9', 'T', 'J', 'Q', 'K', 'A']
JOKER_DECK = ['J', '2', '3', '4', '5', '6', '7', '8', '9', 'T', 'Q', 'K', 'A']
    

# using decorator for sorting
@total_ordering
class Hand(object):
    def __init__(self, cards, bid, deck=DEFAULT_DECK):
        self.cards = cards
        self.bid = bid
        self.card_values = deck
        self.power = 0

    def __str__(self):
        return f'cards: {self.cards}, bid: {self.bid}, power: {self.power}'

    def __lt__(self, obj):
        if self.power < obj.power:
            return True
        elif self.power == obj.power:
            for i in range(len(self.cards)):
                if self.card_values.index(self.cards[i]) < self.card_values.index(obj.cards[i]):
                    return True
                elif self.card_values.index(self.cards[i]) > self.card_values.index(obj.cards[i]):
                    return False
        return False

    def __eq__(self, obj):
        if self.power != obj.power:
            return False
        return self.cards == obj.cards



class HandDefaultDeck(Hand):
    def __init__ (self, cards, bid):
        super().__init__(cards, bid, deck=DEFAULT_DECK)
        self._power()

    def _power(self):
        counter = dict()
        for i in range(1, 4):
            counter[i] = set()
        # man this is ugly :)
        for label in self.cards:
            c = self.cards.count(label) 
            if c == 5:
                self.power = FIVE
                return
            elif c == 4:
                self.power = FOUR
                return
            else:
                counter[c].add(label)
        # end even uglier :)
        if len(counter[2]) > 0 and len(counter[3]) > 0:
            self.power = FHOUSE
        elif len(counter[3]) > 0:
            self.power = THREE
        elif len(counter[2]) > 1:
            self.power = TPAIR
        elif len(counter[2]) > 0:
            self.power = PAIR
        else:
            self.power = HIGH


class HandJokerDeck(Hand):
    def __init__ (self, cards, bid):
        super().__init__(cards, bid, deck=JOKER_DECK)
        self._power()

    def _power(self):
        counter = dict()
        for i in range(1, 5):
            counter[i] = set()

        jokers_count = 0
        # man this is ugly :)
        for label in self.cards:
            if label == 'J':
                jokers_count += 1
                continue
            c = self.cards.count(label) 
            if c == 5:
                # there are five cards of same label, 
                # no jokers, good
                self.power = FIVE
                return
            # otherwise we'll just memorize the counts
            else:
                counter[c].add(label)

        if len(counter[4]) > 0:
            self.power = FOUR
        if len(counter[2]) > 0 and len(counter[3]) > 0:
            self.power = FHOUSE
        elif len(counter[3]) > 0:
            self.power = THREE
        elif len(counter[2]) > 1:
            self.power = TPAIR
        elif len(counter[2]) > 0:
            self.power = PAIR
        else:
            self.power = HIGH

        # let's apply jokers
        for _ in range(jokers_count):
            if self.power == PAIR:
                self.power = THREE 
            elif self.power == TPAIR:
                self.power = FHOUSE
            elif self.power == THREE:
                self.power = FOUR 
            else:
                self.power += 1

            if self.power == FIVE:
                break
        
  
class HandsSet(object):
    def __init__(self, inp, hand_subclass=HandDefaultDeck):
        self.hands = []
        for line in inp:
            s = line.split()
            h = hand_subclass(s[0], int(s[1]))
            self.hands.append(h)

    def __str__(self):
        return '\n'.join(str(h) for h in self.hands)

    def sort(self):
        self.hands = sorted(self.hands)

    def get_total_winnings(self):
        self.sort()
        total = 0
        for rank, hand in enumerate(self.hands, start=1):
            total += (rank * hand.bid)
        return total


if __name__ == '__main__':
    r = tools.Reader('test_input.txt')
    lines = r.read()

    hs = HandsSet(lines)

    x = hs.get_total_winnings()
    print(x)

    hs = HandsSet(lines, hand_subclass=HandJokerDeck)
    hs.sort()
    #print(hs)

    y = hs.get_total_winnings()
    print(y)

    r = tools.Reader('input.txt')
    lines = r.read()

    hs = HandsSet(lines)

    x = hs.get_total_winnings()
    print(x)

    hs = HandsSet(lines, hand_subclass=HandJokerDeck)
    hs.sort()
    #print(hs)

    y = hs.get_total_winnings()
    print(y)