Check “five in a row” for a 15x15 chess board (part 2)

Question

This is a follow-up from here. I have a further optimization for space complexity by store only direction status data for one row for each directions (refer to variable stats). Any advice on algorithm time complexity improvement, code bugs or code style advice is highly appreciated.

Problem:

For a given board, judge who had complete a continuously line (could be horizontal, vertical or diagonal) of 5 pieces. Input is a 15×15 chess board, could be white (1), black (2) or empty (0). Return if white win, if black win, if both win or if none of them win.

import random
import pprint
def check_five(matrix):
    directions = [(0,-1),(-1,0),(-1,-1),(-1,1)] # (row, col)
    stats = [[0] * len(directions) for _ in range(len(matrix[0]))]
    WHITE = 1
    BLACK = 2
    EMPTY = 0
    WIN = 5
    BLACK_WIN = False
    WHITE_WIN = False
    # init first row
    for i in range(len(matrix[0])):
        if matrix[0][i] != EMPTY:
            for j in range(len(directions)):
                stats[i][j] = 1
    for row in range(1, len(matrix)):
        for col in range(0, len(matrix[0])):
            for d_i, d_v in enumerate(directions):
                if matrix[row][col] == EMPTY:
                    stats[col][d_i] = 0
                elif len(matrix)>row+d_v[0] >= 0 and len(matrix[0])>col+d_v[1] >= 0 and matrix[row][col] == matrix[row+d_v[0]][col+d_v[1]]:
                    stats[col][d_i] += 1
                else:
                    stats[col][d_i] = 1
                if stats[col][d_i] == WIN:
                    if matrix[row][col] == WHITE:
                        WHITE_WIN = True
                    else:
                        BLACK_WIN = True
                    if WHITE_WIN and BLACK_WIN:
                        return (True, True)

    return (WHITE_WIN, BLACK_WIN)

if __name__ == "__main__":
    matrix = [[random.choice([0,1,2]) for _ in range(15)] for _ in range(15)]
    pprint.pprint(matrix, indent=4)
    print check_five(matrix)

Answer 1

What an interesting little problem!

You don't need to accumulate any stats at all, if you just test each cell for whether it is at the centre of a line of 5 all the same colour. This also means that, for example, you don't need to look at first two and last two columns when searching for a horizontal line, and similarly for rows and diagonals.

Here's a (rather wholesale) modification of your code. All the real work is done in the winning_cell() function, which just tests the four lines that this cell is at the centre of. This leaves the main loop in check_five() very clean.

Although there's a lot of work in winning_cell, on a randomised matrix this runs about 30% faster.

(EDIT: this is in Python 3.6 - I've just noticed you are on 2.7. Apart from the brackets round the print statement I'm not sure what else this would affect.)

import random
import pprint

# Offsets from the current cell to cells that might form five in a row
# with the current cell as centre, calculated here so as to be outside
# the main loop
OFFSETS = (-2,-1,+1,+2)
DIAG_1 = [*zip(OFFSETS,OFFSETS)]
DIAG_2 = [*zip(OFFSETS,reversed(OFFSETS))]

def winning_cell(m, row, col):
    """ True if the given cell is at the centre of a row of 5 """
    # Note that when checking for a horizontal line we don't need to
    # check the first two or last two columns - the same logic applies
    # to rows and diagonals.
    colour = m[row][col]
    # check within row
    if (1 < col < len(m[0]) - 2 and
        all([m[row][col + c] == colour for c in OFFSETS])):
        return True
    # check within column
    if (1 < row < len(m) - 2 and
        all([m[row+r][col] == colour for r in OFFSETS])):
        return True
    #check diagonals
    if (1 < col < len(m[0]) - 2 and
        1 < row < len(m) - 2 and
        (all([m[row + r][col + c] == colour for r, c in DIAG_1]) or
         all([m[row + r][col + c] == colour for r, c in DIAG_2]))):
        return True
    return False

def check_five(matrix):
    WHITE = 1
    BLACK = 2
    EMPTY = 0
    BLACK_WIN = False
    WHITE_WIN = False

    # main loop to check cells
    for row in range(len(matrix)):
        for col in range(len(matrix[0])):
            cell = matrix[row][col]
            if (cell == EMPTY or
                cell == WHITE and WHITE_WIN or
                cell == BLACK and BLACK_WIN):
                pass
            else:
                if (cell == WHITE and winning_cell(matrix, row, col)):
                    WHITE_WIN = True
                if (cell == BLACK and winning_cell(matrix, row, col)):
                    BLACK_WIN = True
                if WHITE_WIN and BLACK_WIN:
                    return (True, True)

    return (WHITE_WIN, BLACK_WIN)


if __name__ == "__main__":
    matrix = [[random.choice([0,1,2]) for _ in range(15)] for _ in range(15)]
    pprint.pprint(matrix, indent=4)
    print (check_five(matrix))