Connect K

"""
             R
             B
       B R B R
 ... _ _ _ _ _ ... 
    -1 0 1 2 3    

k = 3
Player R inserts at 2
             R
       B   B B
       B R R R
 ... _ _ _ _ _ ... 
    -1 0 1 2 3    
"""


class Token:
    def __init__(self, color):
        self.color = color

    def __eq__(self, other):
        self.color == other.color

    def __hash__(self):
        return hash(self.color)

from collections import defaultdict
class Board:
    def __init__(self, k):
        self.columns = defaultdict(list)
        self.k = k

    def _check_count(placed_token, gen):
        count = 0
        for new_token in gen():
            if new_token != placed_token:
                return count
            count += 1
        return count
    
    def _horizontal_generator(placed_token, x, y, x_direction, y_direction):
        col = self.columns[x]
        for delta in range(1, self.k):
            new_x = x + delta * x_direction
            new_y = y + delta * y_direction
            new_column = self.columns[new_x]
            if len(new_column) < new_y + 1: # since we are 1 indexed
                return
            yield new_column[-new_y - 1]

    def _won_horizontally(self, token, column):
        col = self.columns[column]
        def _horizontally(placed_token, y, direction):
            # direciton is -1 or 1, -1 is left, 1 is right
            # y is 0-indexed
            count = 0
            for delta in range(1, self.k):
                new_column = self.columns[column + delta * direction]
                if len(new_column) < y + 1: # since we are 1 indexed
                    return count
                # we want to count backwards from the back, python -1 == last element => y = 0
                if new_column[-y - 1] != placed_token:
                    return count
                # they are the same
                count += 1
            return count # all are the same  
        def _horizontally_left(placed_token, y) -> int:
            return _horizontally(placeD_token, y, -1)        
        def _horizontally_right(placed_token, y) -> int:
            return _horizontally(placeD_token, y, 1)   
        wins = set()
        for y, placed_token in enumerate(col[::-1]):
            # 1 for self
            if 1 + _horizontally_left(placed_token, y) + _horizontally_right(placed_token, y) >= self.k:
                wins.add(placed_token)
        return wins

    def _won_vertically(self, token, column) -> List[Token]:
        col = self.columns[column]
        count = 0
        for placed_token in col[-1:]:
            if placed_token != token:
                return set() # there are no winners
            count += 1
            if count == self.k:
                return {placed_token}
        return set()

    def play(self, token, column):
        # return the token type which indicates the identity of the player
        self.columns[column].append(token)
        total_wins = set()
        for win_condition in [self._won_horizontally(token, column), self._won_vertically(token, column)]:
            total_wins |= win_condition()
        return total_wins