powerfour.py

powerfour.py

# Power 4
# two players
# one grid: 7 columns, 6 rows
# turn by turn
# each turn: put a coin at the top of a column,
# it falls and reaches the lowest available row in the column

# display ~

from dataclasses import dataclass


@dataclass
class Grid:
    # 2d array
    # None if no coin
    # id of player if coin
    # lowest row is index == 0 in axis=1
    state: list[list[str | None]]


@dataclass
class Action:
    player: str
    column: int


N_COLUMNS = 7
N_ROWS = 6

PLAYER_A = "A"
PLAYER_B = "B"


def check_completed(
    grid: Grid,
    last_column: int,
    last_row: int,
    current_player: str,
) -> bool:
    vertical = grid.state[last_column][max(0, last_row - 4) : min(N_ROWS, last_row + 4)]
    for i_window in range(len(vertical) - 4):
        window = vertical[i_window : i_window + 4]
        if all(x == current_player for x in window):
            return True

    min_i_diagonal = next(
        i for i in range(-4, 0 + 1) if last_row - i >= 0 and last_column - i >= 0
    )
    max_i_diagonal = next(
        i
        for i in range(4, 0 - 1, -1)
        if last_row + i < N_ROWS and last_column + i < N_COLUMNS
    )
    diagonal = [
        grid.state[last_column + i][last_row + i]
        for i in range(min_i_diagonal, max_i_diagonal + 1)
    ]
    for i_window in range(len(diagonal) - 4):
        window = diagonal[i_window : i_window + 4]
        if all(x == current_player for x in window):
            return True
    return False


def game(inputs: list[int]):
    """
    inputs: list of columns where coin is inserted each turn
    """
    grid = Grid(state=[[None for _ in range(N_ROWS)] for _ in range(N_COLUMNS)])
    completed: bool = False
    current_player: str = PLAYER_A
    i_turn: int = 0
    winner: str

    # game loop
    while not completed:
        if i_turn >= len(inputs):
            raise Exception("Not enough inputs to complete")
        action = Action(column=inputs[i_turn], player=current_player)

        lowest = min(
            i_row
            for (i_row, cell) in enumerate(grid.state[action.column])
            if cell is None
        )
        grid.state[action.column][lowest] = action.player

        i_turn += 1
        completed = check_completed(
            grid=grid,
            last_column=action.column,
            last_row=lowest,
            current_player=current_player,
        )
        if completed:
            winner = current_player

        if current_player == PLAYER_A:
            current_player = PLAYER_B
        else:
            current_player = PLAYER_A

    t_grid = [[grid.state[c][r] for c in range(N_COLUMNS)] for r in range(N_ROWS)]
    print()
    for line in reversed(t_grid):
        print(line)
    print()

    print(winner)

    return winner


def test1():
    assert game([0, 1, 0, 2, 0, 2, 0]) == PLAYER_A


def test2():
    assert game([1, 0, 2, 0, 3, 0, 4, 0]) == PLAYER_B


def test3():
    assert game([0, 1, 1, 2, 2, 3, 2, 3, 3, 4, 3]) == PLAYER_A


test1()
test2()
test3()