Connect Four board with gravity drops and board snapshots

Clarify first. Ask four things before writing code. Which way is row 0? Use the printed convention: row 0 is the top, so a piece settles at the largest free row index. What happens on a full column? The board stays the same and you still return a snapshot. What happens on a column outside 0..cols-1? Raise an error, because that is a caller bug and not a game event. Are colors fixed to two players? No. Any single character except the empty marker '0' is a valid color, so do not hard-code red and yellow.

Brute force. On each drop, scan the column from the bottom row upward and stop at the first '0'. That costs O(rows) per drop. It is correct and is a fine first answer to say out loud.

Optimal: keep a height per column. Store heights[c], the number of pieces already in column c. The landing row is rows - 1 - heights[c], so a drop is O(1) before the copy. The column is full exactly when heights[c] == rows.

Return an independent copy. The caller gets a snapshot, so changes they make to it must not reach the board, and later drops must not change a snapshot they hold. Copy each row ([row[:] for row in grid]). grid[:] or list(grid) is the classic bug: it copies the outer list but shares the inner rows. The copy costs O(rows * cols), and that dominates the drop. Say this, and offer a cheaper read API (cell(r, c)) if the interviewer cares about cost.

Part 2: did that drop win? Only the piece just placed can create a new line, so check from it alone. For each of the four directions (horizontal, vertical, and the two diagonals), count matching pieces walking forward and walking backward from the new piece, and add 1 for the piece itself. A total of at least k (4 by default) is a win. This costs O(k) per drop instead of O(rows * cols) for a full board scan.

Invariants.

• For every column c, cells rows-1 down to rows-heights[c] hold pieces, and every cell above them is '0'. Gravity is never violated.
• 0 <= heights[c] <= rows.
• No snapshot shares a list with the internal grid.

from __future__ import annotations

from typing import List, Optional, Tuple

EMPTY = "0"
DIRECTIONS: Tuple[Tuple[int, int], ...] = ((0, 1), (1, 0), (1, 1), (1, -1))


class ConnectFourBoard:
    def __init__(self, rows: int, cols: int, k: int = 4) -> None:
        if rows <= 0 or cols <= 0:
            raise ValueError("rows and cols must be positive")
        self.rows = rows
        self.cols = cols
        self.k = k
        self._grid: List[List[str]] = [[EMPTY] * cols for _ in range(rows)]
        self._heights: List[int] = [0] * cols
        self._last: Optional[Tuple[int, int]] = None  # (row, col) of the last placed piece

    def drop(self, color: str, col: int) -> List[List[str]]:
        """Drop a piece and return an independent copy of the board."""
        if len(color) != 1 or color == EMPTY:
            raise ValueError(f"invalid color: {color!r}")
        if not 0 <= col < self.cols:
            raise ValueError(f"column out of range: {col}")
        if self._heights[col] < self.rows:  # a full column leaves the board unchanged
            row = self.rows - 1 - self._heights[col]
            self._grid[row][col] = color
            self._heights[col] += 1
            self._last = (row, col)
        else:
            self._last = None
        return self.snapshot()

    def snapshot(self) -> List[List[str]]:
        return [row[:] for row in self._grid]  # deep copy of each row, not grid[:]

    def is_full(self, col: int) -> bool:
        return self._heights[col] == self.rows

    def board_full(self) -> bool:
        return all(h == self.rows for h in self._heights)

    # Part 2: does the last successful drop make k in a row?
    def last_move_wins(self) -> bool:
        if self._last is None:
            return False
        r, c = self._last
        color = self._grid[r][c]
        for dr, dc in DIRECTIONS:
            count = 1 + self._run(r, c, dr, dc, color) + self._run(r, c, -dr, -dc, color)
            if count >= self.k:
                return True
        return False

    def _run(self, r: int, c: int, dr: int, dc: int, color: str) -> int:
        n = 0
        r, c = r + dr, c + dc
        while 0 <= r < self.rows and 0 <= c < self.cols and self._grid[r][c] == color:
            n += 1
            r, c = r + dr, c + dc
        return n

• __init__: O(rows * cols) time and space for the grid, plus O(cols) for heights.
• drop: O(1) to place the piece, plus O(rows * cols) time and space for the returned copy. The brute-force column scan would add O(rows), which the copy hides anyway.
• snapshot: O(rows * cols) time and space.
• is_full: O(1). board_full: O(cols).
• last_move_wins: O(k) time when runs are capped at k - 1, O(rows + cols) as written, and O(1) extra space.

• b = ConnectFourBoard(3, 4), then b.drop('R', 1) returns [['0','0','0','0'], ['0','0','0','0'], ['0','R','0','0']].
• Then b.drop('Y', 1) puts 'Y' at row 1, column 1, above the 'R'.
• Fill a column: on a 2-row board, drop('R',0), drop('Y',0), then drop('B',0) returns the same board as before, and last_move_wins() is False.
• Snapshot independence: s = b.drop('R', 0), then s[2][0] = 'X'. The next b.snapshot()[2][0] is still 'R'. A snapshot taken earlier does not change after later drops.
• A 1x1 board: the first drop fills it. The second drop leaves it unchanged, and board_full() is True.
• drop('R', 4) on a 4-column board raises ValueError. So do drop('0', 0) and drop('RR', 0).
• Win checks on a 6x7 board with k = 4:
  • Horizontal: R in columns 0 through 3 on the bottom row, with Y moves between them stacked elsewhere. The fourth R wins.
  • Vertical: four R pieces in column 2 win.
  • Diagonal: R lands at (5,0), (4,1), (3,2), (2,3) on supporting pieces, and the last one wins. Check the anti-diagonal the same way.
  • A gap breaks a line: R R Y R on the bottom row is not a win.
  • The winning piece in the middle of a line counts both sides: R at columns 0, 1, 3, then R at column 2 wins.

• Turn order and game state. Make the game enforce alternating players, reject moves after a win or a draw, and report IN_PROGRESS, WON(color), or DRAW. Put this in a Game class that wraps the board, so the board stays a dumb container with no rules.
• Undo. Keep a stack of moved columns. Undo pops a column, clears grid[rows - heights[c]][c], and decrements heights[c]. That is O(1), with no snapshots to store.
• Cheap reads instead of full copies. If the caller polls the board often, copying it every time is wasteful. Return the changed (row, col, color) as a diff, or an immutable view such as a tuple of strings that is rebuilt only for the changed row. This makes the copy O(cols) instead of O(rows * cols).
• Connect K on a large or unbounded board. Keep the check local to the last move and cap each walk at k - 1. For a very wide sparse board, store pieces in a dict keyed by (row, col) with heights per column. Memory then grows with the number of pieces, not with the board area.
• More than two players, or a bot. The color is already a parameter, so N players work as they are. A simple bot tries every non-full column for an immediate win, then for a block. A stronger one runs minimax with alpha-beta pruning, using drop and undo on one board instead of copying it.

Mid-level (L4)

• Asks which way row 0 points and what happens on a full column, then writes a correct drop.
• Returns a real deep copy and can explain why grid[:] is wrong.
• Tests a full column, an invalid column, and the 1x1 board.

Senior (L5)

• Uses per-column heights for O(1) placement without being asked, and states that the copy dominates the cost.
• Writes the win check from the last move only, covering all four directions and both sides of the new piece.
• Keeps rules (turns, win, draw) out of the board class, and covers each win direction with a test.

Staff (L6+)

• Discusses the API contract: full snapshot versus diff versus immutable view, and who pays for the copy.
• Extends cleanly to undo, connect K, N players, and sparse boards without rewriting the core.
• Explains how the design supports a search bot, with in-place drop and undo and O(k) win checks, and what thread safety would need if several clients shared a board.