connect-4

Connect-4-Style Board Game

Implement a two-player drop-disc board game in the Connect Four family. Players alternate turns; on each move, a disc is dropped into a column and falls to the lowest empty cell in that column. The board dimensions and target line length K should be configurable. After each move, determine whether the moving player has formed K consecutive discs horizontally, vertically, or diagonally.

Requirements

• make_move(player, column) drops the player's disc into the lowest empty cell of the given column.
• Raise an error if the column index is invalid or the column is already full.
• After each move, check whether the moving player has won.
• Support arbitrary board dimensions and arbitrary K.
• Track the game state:
• IN_PROGRESS
• PLAYER_X_WINS
• PLAYER_O_WINS
• DRAW

Design Notes

• Keep the object design focused:
• Board owns the grid and falling-disc logic.
• Game owns turn order, game state, and win detection.
• Player represents the two players.
• Implement Board.drop(column, color) to return the final position (row, col) of the dropped disc.
• Implement Board.winning_run_through(row, col, k) to check whether the disc at (row, col) completes a winning line.
• Implement Game.play(column) to apply one move for the current player and update the game state.
• For win detection, scan from the last move across the four axes:
• horizontal
• vertical
• main diagonal
• anti-diagonal
• For each axis, count consecutive discs of the same player in both directions through the last move. If the total is at least K, the player wins.
• The intended win check is O(K) per move, not O(rows * cols).

Edge Cases

• Invalid column index.
• Full column.
• Move attempted after the game has already ended.
• Board fills with no winner, resulting in a draw.
• K is larger than any possible line on the board. In this case, no player can win, so the game can only end in a draw.

Tests

• Horizontal win.
• Vertical win.
• Diagonal win.
• Full-column rejection.
• Draw.

Follow-Up

If asked how to add an AI player, use minimax with alpha-beta pruning. The AI can simulate future moves and evaluate positions using the same last-move-based O(K) win detection.

---

解法

from enum import Enum
from typing import Callable, List, Tuple

class Player(Enum):
    X = "X"
    O = "O"

    def other_player(self):
        if self == Player.X:
            return Player.O
        return Player.X

class Status(Enum):
    IN_PROGRESS = "IN_PROGRESS"
    PLAYER_X_WINS = "PLAYER_X_WINS"
    PLAYER_O_WINS = "PLAYER_O_WINS"
    DRAW = "DRAW"


class Board:
    def __init__(self, rows: int, cols: int):
        self.rows = rows
        self.cols = cols
        self.index = [0 for i in range(cols)]
        # self.players[i][j]: i-th col, j-th row
        self.players = [[None for i in range(rows)] for j in range(cols)]
        self.spaces = rows * cols

    def play(self, col: int, player: Player) -> Tuple[int, int]:
        if col < 0 or col >= self.cols:
            raise ValueError("invalid column index")
        if self.index[col] == self.rows:
            raise ValueError(f"column {col} has no space left")
        self.players[col][self.index[col]] = player
        self.index[col] += 1
        self.spaces -= 1
        return col, self.index[col] - 1

    def is_valid(self, col: int, row: int) -> bool:
        if col < 0 or col >= self.cols:
            return False
        if row < 0 or row >= self.rows:
            return False
        return True

    def check_one_direction(self, col: int, row: int, dir: Tuple[int, int]) -> int:
        number = 1
        player = self.players[col][row]
        c = col
        r = row
        while True:
            c += dir[0]
            r += dir[1]
            if self.is_valid(c, r) and self.players[c][r] == player:
                number += 1
            else:
                break
        c = col
        r = row
        while True:
            c -= dir[0]
            r -= dir[1]
            if self.is_valid(c, r) and self.players[c][r] == player:
                number += 1
            else:
                break
        return number

    def check(self, col: int, row: int, k: int) -> bool:
        directions = [(1, 0), (0, 1), (1, 1), (1, -1)]
        return any([self.check_one_direction(col, row, direction) >= k for direction in directions])


class Game:
    def __init__(self, rows: int, cols: int, k: int):
        self.board = Board(rows, cols)
        self.k = k
        self.current_player = Player.X
        self.status = Status.IN_PROGRESS

    def play(self, column: int) -> Status:
        if self.status != Status.IN_PROGRESS:
            raise RuntimeError("Game Ended!")
        col, row = self.board.play(column, self.current_player)
        if self.board.check(col, row, self.k):
            if self.current_player == Player.X:
                self.status = Status.PLAYER_X_WINS
            else:
                self.status = Status.PLAYER_O_WINS
        elif self.board.spaces == 0:
            self.status = Status.DRAW
        self.current_player = self.current_player.other_player()
        return self.status


def check(rows: int, cols: int, k: int, actions: List[int], desc: str, expected: Status) -> None:
    g = Game(rows, cols, k)
    status = Status.IN_PROGRESS
    for action in actions:
        status = g.play(action)
    print(f"{status == expected}, desc: {desc}, actual: {status}, expected: {expected}")


def error_name(fn: Callable[[], None]) -> str:
    try:
        fn()
    except Exception as error:
        return type(error).__name__
    return "NO_ERROR"


def check_error(desc: str, fn: Callable[[], None], expected: str) -> None:
    actual = error_name(fn)
    print(f"{actual == expected}, desc: {desc}, actual: {actual}, expected: {expected}")


def main():
    check(
        6,
        7,
        4,
        [0, 0, 1, 1, 2, 2, 3],
        "X_horizontal_wins",
        Status.PLAYER_X_WINS,
    )
    check(
        6,
        7,
        4,
        [0, 1, 0, 1, 0, 1, 0],
        "X_vertical_wins",
        Status.PLAYER_X_WINS,
    )
    check(
        6,
        7,
        4,
        [0, 1, 1, 2, 2, 3, 2, 3, 5, 3, 3],
        "X_main_diagonal_wins",
        Status.PLAYER_X_WINS,
    )
    check(
        6,
        7,
        4,
        [3, 2, 2, 1, 6, 1, 1, 0, 6, 0, 6, 0, 0],
        "X_anti_diagonal_wins",
        Status.PLAYER_X_WINS,
    )
    check(
        2,
        2,
        3,
        [0, 1, 0, 1],
        "draw_when_k_is_larger_than_possible_line",
        Status.DRAW,
    )

    check_error("invalid_column_rejected", lambda: Game(6, 7, 4).play(7), "ValueError")

    full_column_game = Game(2, 2, 3)
    full_column_game.play(0)
    full_column_game.play(0)
    check_error("full_column_rejected", lambda: full_column_game.play(0), "ValueError")

    ended_game = Game(1, 1, 1)
    ended_game.play(0)
    check_error("move_after_game_ended_rejected", lambda: ended_game.play(0), "RuntimeError")

if __name__ == "__main__":
    main()