Generalized tic-tac-toe with k in a row

Question: Generalized tic-tac-toe with k in a row

Companies: Databricks, Hebbia

Difficulty: Medium

Problem statement

Design a two-player Tic-Tac-Toe game on a board with n rows and m columns. A player wins by placing k of their own marks consecutively in one line: horizontal, vertical, diagonal, or anti-diagonal.

Implement the class TicTacToe:

• TicTacToe(n: int, m: int, k: int) creates an empty n x m board with win length k (1 <= k <= max(n, m)).
• move(row: int, col: int, player: int) -> int places a mark for player (1 or 2) on the empty cell (row, col) and returns:
  • 0 if nobody has won yet,
  • 1 or 2 if this move wins the game for that player.

Rules:

1. Every move is on an empty, in-bounds cell. Validating this is still good practice.
2. The game ends at the first win. After that, no more moves are allowed.
3. The marks must be adjacent. A row with k marks of one player that has a gap is not a win.

Unlike the classic version, k may be smaller than the board's sides, so a win does not have to fill a whole row, column or diagonal. Aim for a move that is faster than rescanning the board.

Constraints: 1 <= n, m <= 1000, 1 <= k <= max(n, m), at most n * m calls to move.

Example 1

game = TicTacToe(4, 5, 3)
game.move(3, 0, 1)  # 0
game.move(0, 4, 2)  # 0
game.move(2, 1, 1)  # 0
game.move(0, 3, 2)  # 0
game.move(1, 2, 1)  # 1  -> (3,0), (2,1), (1,2) form an anti-diagonal of length 3

Example 2

game = TicTacToe(1, 5, 3)
game.move(0, 0, 1)  # 0
game.move(0, 1, 1)  # 0
game.move(0, 3, 1)  # 0  -> three marks in the row, but not consecutive
game.move(0, 2, 1)  # 1  -> columns 0..3 are now one run of 4 >= 3

(Player turns are not enforced in Example 2, to keep it short.)

Follow-up: Can you report a draw when the board fills up without a winner?

Approach

This generalizes LeetCode 348 (Design Tic-Tac-Toe). The well-known O(1) trick
there, one signed counter per row, column and diagonal, only works when a win
means filling a whole line (k == n == m). With k smaller than the board,
a row can hold k marks for a player without them being next to each other,
so counters can't tell a win from a near miss. We need a different invariant.

Brute force. After each move, scan the whole board for any run of k
equal marks in the four directions. That's O(n·m·k) per move and does work
that has nothing to do with the move just made.

Key observation. Before this move nobody had won, since play stops at the
first win. So if a win exists now, it must use the cell just played. We only
have to look at lines through (r, c).

Optimal method: count outward from the move. There are four line
directions: horizontal (0, 1), vertical (1, 0), diagonal (1, 1) and
anti-diagonal (1, -1). For each one:

1. Walk forward from (r, c) while the cell is in bounds and holds the
   player's mark, counting cells.
2. Walk backward the same way along the opposite direction.
3. The run through the move is 1 + forward + backward. If it's >= k, the
   player wins.

Each walk stops after k - 1 steps at most, because a longer run can't
change the answer. So a move costs O(k), not O(n·m).

Invariants.

• winner is 0 until the first winning move, and after that it never
  changes. Any later move raises an error, as the rules say.
• Every non-empty cell holds 1 or 2, and before the winning move no run of
  k equal marks exists anywhere on the board.
• moves_made counts filled cells. When it reaches n * m and nobody has won,
  the game is a draw.

Code

from enum import IntEnum


class Result(IntEnum):
    IN_PROGRESS = 0
    PLAYER_1 = 1
    PLAYER_2 = 2
    DRAW = 3


class GameOverError(Exception):
    pass


class TicTacToe:
    """Two-player k-in-a-row on an n x m board (n rows, m columns)."""

    DIRECTIONS = ((0, 1), (1, 0), (1, 1), (1, -1))

    def __init__(self, n: int, m: int, k: int) -> None:
        if n <= 0 or m <= 0 or k <= 0:
            raise ValueError("n, m and k must be positive")
        if k > max(n, m):
            raise ValueError("k cannot exceed the longest side of the board")
        self.n, self.m, self.k = n, m, k
        self.board: list[list[int]] = [[0] * m for _ in range(n)]
        self.moves_made = 0
        self.result = Result.IN_PROGRESS

    def move(self, row: int, col: int, player: int) -> int:
        """Place player's mark. Returns 0 (no winner yet), 1 or 2 (winner),
        or 3 (the board is full with no winner)."""
        if self.result != Result.IN_PROGRESS:
            raise GameOverError("the game is already over")
        if player not in (1, 2):
            raise ValueError("player must be 1 or 2")
        if not self._in_bounds(row, col) or self.board[row][col] != 0:
            raise ValueError(f"illegal move at ({row}, {col})")

        self.board[row][col] = player
        self.moves_made += 1

        if self._wins_through(row, col, player):
            self.result = Result(player)
        elif self.moves_made == self.n * self.m:
            self.result = Result.DRAW
        return int(self.result)

    def _wins_through(self, row: int, col: int, player: int) -> bool:
        for dr, dc in self.DIRECTIONS:
            run = 1
            run += self._count(row, col, dr, dc, player)
            run += self._count(row, col, -dr, -dc, player)
            if run >= self.k:
                return True
        return False

    def _count(self, row: int, col: int, dr: int, dc: int, player: int) -> int:
        """Consecutive marks of player after (row, col) along (dr, dc),
        capped at k - 1 because a longer run cannot change the answer."""
        count = 0
        r, c = row + dr, col + dc
        while (count < self.k - 1 and self._in_bounds(r, c)
               and self.board[r][c] == player):
            count += 1
            r, c = r + dr, c + dc
        return count

    def _in_bounds(self, row: int, col: int) -> bool:
        return 0 <= row < self.n and 0 <= col < self.m

Complexity

• __init__: O(n·m) time and space for the board.
• move: O(k) time. There are four directions and two walks of at most
  k - 1 cells in each, so at most 8(k - 1) cells are read. It uses O(1)
  extra space.
• In total, a full game of n·m moves costs O(n·m·k), compared with
  O((n·m)²·k) when the board is rescanned after every move.

Test cases

• n=3, m=3, k=3, moves (0,0,1) (1,0,2) (0,1,1) (1,1,2) (0,2,1): the last
  move returns 1 (top row).
• n=4, m=5, k=3, moves (3,0,1) (0,4,2) (2,1,1) (0,3,2) (1,2,1): returns 1
  on the last move, an anti-diagonal of length 3 that doesn't touch a corner.
• Gap in the middle: n=1, m=5, k=3, player 1 at columns 0, 1, 3 gives
  0, 0, 0. Then column 2 returns 1, because the run is 0..3 even though
  the move is in the middle of it.
• Rectangular, vertical: n=5, m=2, k=4, player 2 at (1,1) (2,1) (3,1) (4,1)
  with player 1 moves in between: the fourth player-2 move returns 2.
• k=1: the first move returns that player's number.
• Draw: n=3, m=3, k=3, sequence X O X / X O O / O X X filled with no
  win. The ninth move returns 3.
• A move after a win raises GameOverError. A move on an occupied cell or
  out of bounds raises ValueError.
• k > max(n, m) is rejected in the constructor, because no win would be
  possible.

Follow-ups

1. Process a list of moves and report the winner after each one. Wrap
   move in a loop and collect the results. Many interviewers also want to
   know at which move index the game ended. Stop at the first nonzero
   result, or report later moves as invalid.
2. Add a computer player (random first, then "smarter"). Keep a set of
   empty cells so a random move is an O(1) pick. A greedy bot plays a winning
   cell if it has one, blocks the opponent's winning cell next, and
   otherwise plays randomly. Testing a cell with _wins_through on a
   temporary board costs O(k), so one bot turn is O(n·m·k). Full minimax
   works only on tiny boards. Mention alpha-beta pruning and a depth limit
   with a heuristic that counts open runs.
3. Undo / replay. Keep a stack of (row, col, player). Undo clears the
   cell, decrements moves_made, and sets result back to IN_PROGRESS.
   This is safe because the invariant says that no win existed before the
   last move.
4. More than two players, or an unbounded board (Gomoku on an infinite
   grid). Store marks in a dict[(row, col)] -> player instead of a matrix
   and drop the bounds check. The outward scan doesn't change, and memory
   becomes O(moves).
5. O(1) per move when k equals the board side (the LeetCode 348 case).
   Use per-row, per-column and two diagonal counters (+1 for one player and
   -1 for the other) and check abs(count) == k. Explain why this fails
   when k is smaller: counts ignore adjacency.

Interviewer expectations

Mid-level (L4)

• Gets a correct board and win check working, even if it starts by
  scanning the board, and then moves to checking only lines through the
  last move.
• Handles bounds on both ends of each direction, including the
  anti-diagonal, and counts the run on both sides of the move.
• Gives the O(k) per-move cost and writes a few asserts.

Senior (L5)

• Says at once why the counter trick from LeetCode 348 doesn't work for
  general k, and explains the invariant that a new win must pass
  through the last move.
• Has clean API decisions: a game-over state, a draw result, and clear
  errors for illegal moves. The code is easy to extend to a bot or undo.
• Tests the tricky cases without being asked: a gap filled in the middle,
  rectangular boards, k = 1, and a draw.

Staff (L6+)

• Leads the conversation through the follow-ups: an unbounded or sparse
  board, N players, undo, and a bot strategy with a cost analysis for each.
• Talks about where the game state should live if it's served as a
  multiplayer service (concurrent moves, turn order enforced on the
  server, replaying from a move log).
• Says when to stop: when minimax becomes intractable and what heuristic
  to use instead.