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.
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 = [ * len(directions) for _ in range(len(matrix))] WHITE = 1 BLACK = 2 EMPTY = 0 WIN = 5 BLACK_WIN = False WHITE_WIN = False # init first row for i in range(len(matrix)): if matrix[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)): 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 and len(matrix)>col+d_v >= 0 and matrix[row][col] == matrix[row+d_v][col+d_v]: 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)