Solve n diagonals problem with backtracking

n_diagonals.py

def can_be_extended_to_solution(perm, n, m):
    # Too many non diagonal squares
    if perm.count(2) > m*m - n:
        return False

    new = len(perm) - 1
    # New square has no diagonal
    if perm[new] == 2:
        return True

    # Get the row and column of new square
    new_row = int(new / m)
    new_col = int(new % m)

    # Compare every square to this new square
    for cur in range(new):
        cur_row = int(cur / m)
        cur_col = int(cur % m)

        # No comparison needed for non diagonal
        if perm[cur] == 2:
            continue

        # left-right
        if new_row == cur_row and abs(new_col - cur_col) == 1:
            if perm[new] != perm[cur]:
                return False
        # top-bottom
        elif new_col == cur_col and abs(new_row - cur_row) == 1:
            if perm[new] != perm[cur]:
                return False
        # left-diagonal
        elif (new_col - cur_col == 1 and new_row - cur_row == 1) or (cur_col - new_col == 1 and cur_row - new_row == 1):
            if perm[new] == 0 and perm[cur] == 0:
                return False
        # right-diagonal
        elif (new_col - cur_col == 1 and cur_row - new_row == 1) or (cur_col - new_col == 1 and new_row - cur_row == 1):
            if perm[new] == 1 and perm[cur] == 1:
                return False

    return True


def extend(perm, n, m):
    # 0 is left diagonal
    # 1 is right diagonal
    # 2 is no diagonal
    if perm.count(0) + perm.count(1) == n:
        print_solution(perm, n, m)
        exit()

    for k in range(3):
        perm.append(k)

        if can_be_extended_to_solution(perm, n, m):
            extend(perm, n, m)

        perm.pop()


def print_solution(perm, n, m):
    print('Solution: ')

    for i in range(m):
        line = ''
        for j in range(m):
            index = i*m + j
            if index < len(perm):
                if perm[index] == 0:
                    line += '\\'
                elif perm[index] == 1:
                    line += '/'
                else:
                    line += ' '
            else:
                line += ' '
        print(line)


extend([], 16, 5)