Tic tac toe encoding

encode_tic_tac_toe.py

import math

# Combination encode and decode functions copied/adapted from https://stackoverflow.com/a/3143594
def encode_comb(n, comb):
    k = len(comb)

    if k == 0 or k == n:
        return 0

    j = comb[0]

    if k == 1:
        return j

    comb = [x - 1 for x in comb]
    if j == 0:
        return encode_comb(n - 1, comb[1:])

    return math.comb(n - 1, k - 1) + encode_comb(n - 1, comb)

def decode_comb(n, k, m):
    if k == 0:
        return []

    if k == n:
        return list(range(n))

    if k == 1:
        return [m]

    b = math.comb(n - 1, k - 1)
    if m < b:
        return [0] + [x + 1 for x in decode_comb(n - 1, k - 1, m)]

    return [x + 1 for x in decode_comb(n - 1, k, m - b)]

def offset(count_non_blank):
    total_count = 0
    for i in range(count_non_blank):
        total_count += math.comb(9, i) * math.comb(i, i // 2)

    return total_count

def unoffset(enc_board):
    offset = 0
    for count_non_blank in range(10):
        count = math.comb(9, count_non_blank) * math.comb(count_non_blank, count_non_blank // 2)
        if offset + count > enc_board:
            return offset, count_non_blank
        offset += count

    assert(False)

def encode_board(board):
    non_blanks = []
    for i, c in enumerate(board):
        if c != " ":
            non_blanks.append(i)

    count_non_blank = len(non_blanks)
    enc_non_blank = encode_comb(9, non_blanks)

    os = []
    for i, j in enumerate(non_blanks):
        if board[j] == "O":
            os.append(i)

    count_o = len(os)
    enc_o = encode_comb(count_non_blank, os)
    comb_o = math.comb(count_non_blank, count_o)

    enc_board = enc_non_blank * comb_o + enc_o

    return offset(count_non_blank) + enc_board

def decode_board(enc_board):
    offset, count_non_blank = unoffset(enc_board)
    enc_board -= offset

    count_o = count_non_blank // 2
    comb_o = math.comb(count_non_blank, count_o)
    enc_o = enc_board % comb_o
    os = decode_comb(count_non_blank, count_o, enc_o)

    enc_non_blank = enc_board // comb_o
    non_blanks = decode_comb(9, count_non_blank, enc_non_blank)

    board = [" "] * 9
    for i in os:
        board[non_blanks[i]] = "O"

    for i in non_blanks:
        if board[i] == " ":
            board[i] = "X"

    return board


def set_cell(board_index, cell_index, c):
    board = decode_board(board_index)
    board[cell_index] = c
    return encode_board(board)

def get_cell(board_index, cell_index):
    return decode_board(board_index)[cell_index]

index = encode_board(["X", " ", "X", " ", "X", "O", "O", "X", "O"])
board = decode_board(index)
print(index, board)
new_index = set_cell(index, 0, " ")
print(get_cell(new_index, 2))
print(decode_board(new_index))

print([offset(i) for i in range(10)])

print(encode_board(["X", "X", "X", " ", " ", " ", "O", "O", "O"]))

seen_boards = set()
i = 0
for i in range(6046):
    board = tuple(decode_board(i))
    assert(board not in seen_boards)
    seen_boards.add(board)

    i += 1