Code editing task: minimax_to_alphabeta
Topic: DSA; change kind: perfective

Original code:
import copy
from typing import List, Literal, Optional, Tuple

Player = Literal['X', 'O']
WinStatus = Literal[Player, 'TIE', None]


class ConnectNGame:
    """
    A game of Connect N, of width x height, where N is the number of pieces in a row/column/diagonal to win.
    """

    def __init__(self, width, height, n):
        self.width = width
        self.height = height
        self.n = n
        self.board = [[' ' for _ in range(width)] for _ in range(height)]

    def __str__(self):
        return '\n'.join(['|' + '|'.join(row) + '|' for row in self.board])

    def drop(self, column, player: Player) -> bool:
        if column < 0 or column >= self.width:
            return False
        for row in range(self.height - 1, -1, -1):
            if self.board[row][column] == ' ':
                self.board[row][column] = player
                return True
        return False

    def is_won(self) -> WinStatus:
        # Check rows
        for row in self.board:
            for i in range(self.width - self.n + 1):
                if row[i] != ' ' and all(row[i] == row[j] for j in range(i + 1, i + self.n)):
                    return row[i]

        # Check columns
        for j in range(self.width):
            for i in range(self.height - self.n + 1):
                if self.board[i][j] != ' ' and all(self.board[i][j] == self.board[k][j] for k in range(i + 1, i + self.n)):
                    return self.board[i][j]

        # Check diagonals
        for i in range(self.height - self.n + 1):
            for j in range(self.width - self.n + 1):
                if self.board[i][j] != ' ' and all(self.board[i][j] == self.board[i + k][j + k] for k in range(1, self.n)):
                    return self.board[i][j]

        for i in range(self.height - self.n + 1):
            for j in range(self.n - 1, self.width):
                if self.board[i][j] != ' ' and all(self.board[i][j] == self.board[i + k][j - k] for k in range(1, self.n)):
                    return self.board[i][j]

        # Check for tie
        if all(self.board[i][j] != ' ' for i in range(self.height) for j in range(self.width)):
            return 'TIE'

        return None

    def possible_moves(self) -> List[int]:
        return [col for col in range(self.width) if self.board[0][col] == ' ']

    def heuristic(self, player: Player) -> float:
        """
        Returns a heuristic score [-0.9, 0.9] for the board state.
        """
        score = 0

        # center column preference
        center_column = [self.board[i][self.width // 2]
                         for i in range(self.height)]
        center_count = center_column.count(player)
        score += center_count * 0.3

        # check rows, columns, and diagonals for potential wins
        for row in range(self.height):
            for col in range(self.width):
                if self.board[row][col] == ' ':
                    continue

                # horizontal potential
                if col + self.n <= self.width:
                    window = [self.board[row][c]
                              for c in range(col, col + self.n)]
                    score += self.evaluate_window(window, player)

                # vertical potential
                if row + self.n <= self.height:
                    window = [self.board[r][col]
                              for r in range(row, row + self.n)]
                    score += self.evaluate_window(window, player)

                # positive diagonal
                if col + self.n <= self.width and row + self.n <= self.height:
                    window = [self.board[row + i][col + i]
                              for i in range(self.n)]
                    score += self.evaluate_window(window, player)

                # negative diagonal
                if col - self.n >= -1 and row + self.n <= self.height:
                    window = [self.board[row + i][col - i]
                              for i in range(self.n)]
                    score += self.evaluate_window(window, player)

        return score

    def evaluate_window(self, window, player):
        opponent = 'O' if player == 'X' else 'X'
        score = 0
        if window.count(player) == self.n - 1 and window.count(' ') == 1:
            score += 0.5
        if window.count(player) == self.n - 2 and window.count(' ') == 2:
            score += 0.2
        if window.count(opponent) == self.n - 1 and window.count(' ') == 1:
            score -= 0.4
        return score

    def score_position(self, status: WinStatus, player: Player) -> float:
        """
        Assign scores to the board state. 
        Win is 1, loss is -1, tie (or ongoing) is heuristic.
        """
        status = self.is_won()
        if status == player:
            return 1
        elif status == 'TIE':
            return 0
        elif status is None:
            return self.heuristic(player)
        else:
            return -1

    def ai(self, depth: int, maximizing: bool, player: Player) -> Tuple[float, Optional[int]]:
        """ 
        Implements an AI that picks the "best" move using Minimax.
        Returns a tuple of (score, column).
        """
        opponent = 'O' if player == 'X' else 'X'

        if depth == 0:
            return self.score_position(self.is_won(), player), None
        terminal_status = self.is_won()
        if terminal_status is not None:
            return self.score_position(terminal_status, player), None

        moves = self.possible_moves()

        if maximizing:
            max_score = float('-inf')
            best_column = None
            for move in moves:
                temp_game = copy.deepcopy(self)
                temp_game.drop(move, player)
                score, _ = temp_game.ai(depth - 1, False, opponent)
                if score > max_score:
                    max_score = score
                    best_column = move
            return max_score, best_column
        else:
            min_score = float('inf')
            best_column = None
            for move in moves:
                temp_game = copy.deepcopy(self)
                temp_game.drop(move, opponent)
                score, _ = temp_game.ai(depth - 1, True, player)
                if score < min_score:
                    min_score = score
                    best_column = move
            return min_score, best_column

    def best_move(self, player: Player, depth=4) -> int:
        """ Returns the best column for the player using Minimax. """
        _, best_column = self.ai(depth, False, player)
        if best_column is None:
            best_column = self.possible_moves()[0]
        return best_column

Instruction (descriptive):
Augment the minimax algorithm with alpha-beta pruning to make it faster. 
Keep track of an alpha and beta value, which represent the minimum score that the maximizing player is assured of and the maximum score that the minimizing player is assured of respectively.
Utilize these two scores to prune branches of the search tree that cannot possibly contain the optimal move.

Instruction (lazy):
Optimize the AI to find the best move in less steps.

Edited code:
import copy
from typing import List, Literal, Optional, Tuple

Player = Literal['X', 'O']
WinStatus = Literal[Player, 'TIE', None]


class ConnectNGame:
    """
    A game of Connect N, of width x height, where N is the number of pieces in a row/column/diagonal to win.
    """

    def __init__(self, width, height, n):
        self.width = width
        self.height = height
        self.n = n
        self.board = [[' ' for _ in range(width)] for _ in range(height)]

    def __str__(self):
        return '\n'.join(['|' + '|'.join(row) + '|' for row in self.board])

    def drop(self, column, player: Player) -> bool:
        if column < 0 or column >= self.width:
            return False
        for row in range(self.height - 1, -1, -1):
            if self.board[row][column] == ' ':
                self.board[row][column] = player
                return True
        return False

    def is_won(self) -> WinStatus:
        # Check rows
        for row in self.board:
            for i in range(self.width - self.n + 1):
                if row[i] != ' ' and all(row[i] == row[j] for j in range(i + 1, i + self.n)):
                    return row[i]

        # Check columns
        for j in range(self.width):
            for i in range(self.height - self.n + 1):
                if self.board[i][j] != ' ' and all(self.board[i][j] == self.board[k][j] for k in range(i + 1, i + self.n)):
                    return self.board[i][j]

        # Check diagonals
        for i in range(self.height - self.n + 1):
            for j in range(self.width - self.n + 1):
                if self.board[i][j] != ' ' and all(self.board[i][j] == self.board[i + k][j + k] for k in range(1, self.n)):
                    return self.board[i][j]

        for i in range(self.height - self.n + 1):
            for j in range(self.n - 1, self.width):
                if self.board[i][j] != ' ' and all(self.board[i][j] == self.board[i + k][j - k] for k in range(1, self.n)):
                    return self.board[i][j]

        # Check for tie
        if all(self.board[i][j] != ' ' for i in range(self.height) for j in range(self.width)):
            return 'TIE'

        return None

    def possible_moves(self) -> List[int]:
        return [col for col in range(self.width) if self.board[0][col] == ' ']

    def heuristic(self, player: Player) -> float:
        """
        Returns a heuristic score [-0.9, 0.9] for the board state.
        """
        score = 0

        # center column preference
        center_column = [self.board[i][self.width // 2]
                         for i in range(self.height)]
        center_count = center_column.count(player)
        score += center_count * 0.3

        # check rows, columns, and diagonals for potential wins
        for row in range(self.height):
            for col in range(self.width):
                if self.board[row][col] == ' ':
                    continue

                # horizontal potential
                if col + self.n <= self.width:
                    window = [self.board[row][c]
                              for c in range(col, col + self.n)]
                    score += self.evaluate_window(window, player)

                # vertical potential
                if row + self.n <= self.height:
                    window = [self.board[r][col]
                              for r in range(row, row + self.n)]
                    score += self.evaluate_window(window, player)

                # positive diagonal
                if col + self.n <= self.width and row + self.n <= self.height:
                    window = [self.board[row + i][col + i]
                              for i in range(self.n)]
                    score += self.evaluate_window(window, player)

                # negative diagonal
                if col - self.n >= -1 and row + self.n <= self.height:
                    window = [self.board[row + i][col - i]
                              for i in range(self.n)]
                    score += self.evaluate_window(window, player)

        return score

    def evaluate_window(self, window, player):
        opponent = 'O' if player == 'X' else 'X'
        score = 0
        if window.count(player) == self.n - 1 and window.count(' ') == 1:
            score += 0.5
        if window.count(player) == self.n - 2 and window.count(' ') == 2:
            score += 0.2
        if window.count(opponent) == self.n - 1 and window.count(' ') == 1:
            score -= 0.4
        return score

    def score_position(self, status: WinStatus, player: Player) -> float:
        """
        Assign scores to the board state. 
        Win is 1, loss is -1, tie (or ongoing) is heuristic.
        """
        status = self.is_won()
        if status == player:
            return 1
        elif status == 'TIE':
            return 0
        elif status is None:
            return self.heuristic(player)
        else:
            return -1

    def ai(self, depth: int, maximizing: bool, player: Player, alpha: float = float('-inf'), beta: float = float('inf')) -> Tuple[float, Optional[int]]:
        """ 
        Implements an AI that picks the "best" move using Minimax with Alpha-Beta pruning.
        Returns a tuple of (score, column).
        """
        opponent = 'O' if player == 'X' else 'X'

        status = self.is_won()
        if depth == 0 or status is not None:
            return self.score_position(status, player), None

        if maximizing:
            max_score = float('-inf')
            best_column = None
            for move in self.possible_moves():
                temp_game = copy.deepcopy(self)
                temp_game.drop(move, player)
                score, _ = temp_game.ai(
                    depth - 1, False, opponent, alpha, beta)
                if score > max_score:
                    max_score = score
                    best_column = move
                alpha = max(alpha, score)
                if alpha >= beta:
                    break
            return max_score, best_column
        else:
            min_score = float('inf')
            best_column = None
            for move in self.possible_moves():
                temp_game = copy.deepcopy(self)
                temp_game.drop(move, opponent)
                score, _ = temp_game.ai(depth - 1, True, player, alpha, beta)
                if score < min_score:
                    min_score = score
                    best_column = move
                beta = min(beta, score)
                if beta <= alpha:
                    break
            return min_score, best_column

    def best_move(self, player: Player, depth=4) -> int:
        """ Returns the best column for the player using Minimax. """
        _, best_column = self.ai(depth, False, player)
        if best_column is None:
            best_column = self.possible_moves()[0]
        return best_column

Tests:
### START TESTS ###
if True:  # pragma: no cover
    game1 = ConnectNGame(7, 6, 4)
    assert game1.drop(0, 'X')
    assert game1.drop(0, 'O')
    assert game1.drop(0, 'X')
    assert game1.drop(0, 'O')
    assert game1.drop(0, 'X')
    assert game1.drop(0, 'O')
    assert not game1.drop(0, 'X')
    assert not game1.is_won()

    game2 = ConnectNGame(4, 4, 3)
    assert game2.drop(0, 'X')
    assert game2.drop(1, 'X')
    assert game2.drop(2, 'X')
    assert game2.is_won() == 'X'

    game3 = ConnectNGame(4, 4, 3)
    assert game3.drop(0, 'X')
    assert game3.drop(1, 'O')
    assert game3.drop(2, 'X')
    assert game3.drop(3, 'O')
    assert game3.drop(0, 'X')
    assert game3.drop(1, 'O')
    assert game3.drop(2, 'X')

    game4 = ConnectNGame(7, 6, 4)
    assert game4.width == 7
    assert game4.height == 6
    assert game4.n == 4
    assert game4.board == [[' ' for _ in range(7)] for _ in range(6)]
    assert str(game4) == '\n'.join(
        ['|' + '|'.join([' ' for _ in range(7)]) + '|' for _ in range(6)])
    game = ConnectNGame(7, 6, 4)
    assert game.drop(0, 'X') == True
    assert game.drop(0, 'O') == True
    assert game.drop(7, 'X') == False
    assert game.drop(-1, 'O') == False
    # Test for no winner
    game = ConnectNGame(7, 6, 4)
    assert game.is_won() == None

    # Test for a horizontal win
    for col in range(4):
        game.drop(col, 'X')
    assert game.is_won() == 'X'

    # Test for a vertical win
    game = ConnectNGame(7, 6, 4)
    for _ in range(4):
        game.drop(0, 'O')
    assert game.is_won() == 'O'

    # Test for a diagonal win
    game = ConnectNGame(7, 6, 4)
    for i in range(4):
        for j in range(i):
            game.drop(i, 'O')
        game.drop(i, 'X')
    assert game.is_won() == 'X'

    game = ConnectNGame(3, 3, 3)
    for i in range(3):
        for j in range(3):
            player = 'X' if (i + j) % 2 == 0 else 'O'
            game.drop(i, player)
    assert game.is_won() == 'X'
    game = ConnectNGame(3, 3, 4)
    game.board = [['X', 'O', 'X'], ['O', 'X', 'O'], ['O', 'X', 'O']]
    assert game.is_won() == 'TIE'
    assert game.score_position(game.is_won(), 'X') == 0

    game = ConnectNGame(7, 6, 4)
    assert game.possible_moves() == list(range(7))

    game.drop(0, 'X')
    game.drop(0, 'O')
    assert game.possible_moves() == list(range(7))

    for _ in range(6):
        game.drop(1, 'X')
    assert 1 not in game.possible_moves()

    best_move = game.best_move('X', 3)
    assert best_move in range(7)

    game = ConnectNGame(7, 6, 4)
    for i in range(3):
        game.drop(i, 'X')

    best_move_x = game.best_move('X', 1)
    assert best_move_x == 3

    game = ConnectNGame(7, 6, 4)
    for i in range(3):
        game.drop(i, 'O')

    best_move_x = game.best_move('X', 4)
    assert best_move_x == 3


    game = ConnectNGame(7, 6, 4)
    for i in range(3):
        game.drop(i, 'X')
        game.drop(i + 1, 'O')

    best_move_x = game.best_move('X', 4)
    assert best_move_x == 4

    __EVAL_COUNTER = 0  # need a global because of deepcopy

    game = ConnectNGame(7, 6, 4)
    for i in range(2, 5):
        game.drop(i, 'O')

    best_move_x = game.best_move('X', 3)
    assert best_move_x == 1 or best_move_x == 5

    game = ConnectNGame(7, 6, 4)

    game.drop(0, 'X')
    game.drop(1, 'O')
    game.drop(3, 'X')
    game.drop(2, 'O')
    game.drop(4, 'X')
    game.drop(5, 'O')
    game.drop(1, 'X')
    game.drop(0, 'O')
    game.drop(2, 'X')
    game.drop(3, 'O')
    game.drop(2, 'X')
    game.drop(3, 'O')
    game.drop(0, 'X')
    game.drop(3, 'O')
    game.drop(3, 'X')
    game.drop(0, 'X')
    game.drop(1, 'O')
    game.drop(3, 'X')
    game.drop(5, 'O')
    game.drop(1, 'X')
    game.drop(4, 'O')
    game.drop(2, 'X')
    best_move_o = game.best_move('O', 4)
    assert best_move_o == 2
    game.drop(best_move_o, 'O')
    game.drop(4, 'X')
    game.drop(4, 'O')
    game.drop(0, 'X')
    game.drop(4, 'O')
    assert game.best_move('X', 8) == 0


    class __EVAL_ConnectNGameWithCounter(ConnectNGame):
        def __init__(self, width, height, n):
            super().__init__(width, height, n)

        def possible_moves(self):
            global __EVAL_COUNTER
            __EVAL_COUNTER += 1
            return super().possible_moves()

        def reset_counter(self):
            global __EVAL_COUNTER
            __EVAL_COUNTER = 0


    game = __EVAL_ConnectNGameWithCounter(7, 6, 4)
    game.drop(0, 'X')
    game.drop(1, 'O')
    game.drop(3, 'X')

    game.reset_counter()
    _ = game.best_move('X', 4)

    assert __EVAL_COUNTER < 200  # alpha-beta gets 184

    game = __EVAL_ConnectNGameWithCounter(7, 6, 4)
    game.drop(2, 'X')
    game.drop(3, 'O')
    game.drop(2, 'X')

    game.reset_counter()
    _ = game.best_move('X', 4)
    assert __EVAL_COUNTER < 180  # alpha-beta gets 166

    game = __EVAL_ConnectNGameWithCounter(10, 10, 5)

    game.drop(0, 'X')
    game.drop(1, 'O')
    game.drop(3, 'X')
    game.drop(2, 'O')
    game.drop(4, 'X')
    game.drop(5, 'O')
    game.drop(6, 'X')
    game.drop(7, 'O')
    game.drop(8, 'X')
    game.drop(9, 'O')

    game.reset_counter()
    _ = game.best_move('X')
    assert __EVAL_COUNTER < 350  # alpha-beta gets 319

    game = __EVAL_ConnectNGameWithCounter(10, 10, 5)

    game.reset_counter()
    _ = game.best_move('X', 6)  # very deep for a normal minimax
    assert __EVAL_COUNTER < 3500  # alpha-beta gets 3137