Finding number of solutions nqueens 14 by 14 chessboard

Question

I have written a code to solve the N-Queen problem for a 14 by 14 chessboard. There are preplaced queens (see below for the sample input and output). The lines of sample input (after the 3) represent the positions of the preplaced queens. For example, "1 1 2 9 3 6 4 10" refers to the 4 preplaced queens are at (1st row 1st column), (2nd row, 9th column), (3rd row 6th column), (4th row, 10th column).

I am hoping that it will return me the number of solutions. But when I tried running the code, I got time limit exceeded, even for the sample input.

I will try the list of tuples method, but let me know if there is a more efficient implementation. Thanks in advance

[![Question screenshot][1]][1]

Sample Input
3
1 1 2 9 3 6 4 10
1 9 12 8 8 5 2 7
2 10 10 3

Sample Output
39
32
2414

import sys


def nQueens(n, pre_placed_rows, pre_placed_cols, N, state=[], col=1):  
    if col > n: return [state]
    res = []

    for i in range(1, n+1):
        if invalid(state, i): continue
        for sol in nQueens(n, pre_placed_rows, pre_placed_cols, N, state + [i], col+1): 
            if N == 2: # 2 preplaced queens
                if sol[pre_placed_cols[0]-1] == pre_placed_rows[0] and sol[pre_placed_cols[1]-1] == pre_placed_rows[1]: 
                    res.append(sol) #[sol]
            elif N == 3:  # 3 preplaced queens
                if sol[pre_placed_cols[0]-1] == pre_placed_rows[0] and sol[pre_placed_cols[1]-1] == pre_placed_rows[1] and sol[pre_placed_cols[2]-1] == pre_placed_rows[2]:
                    res.append(sol)
            else: # 4 preplaced queens
                if sol[pre_placed_cols[0]-1] == pre_placed_rows[0] and sol[pre_placed_cols[1]-1] == pre_placed_rows[1] and sol[pre_placed_cols[2]-1] == pre_placed_rows[2] and sol[pre_placed_cols[3]-1] == pre_placed_rows[3]:
                    res.append(sol)
    return res

def invalid(s, r2): # is it safe to put the queen 
    if not s: 
        return False
    if r2 in s: 
        return True
    c2 = len(s) + 1
    return any(abs(c1-c2) == abs(r1-r2) for c1, r1 in enumerate(s, 1))

num_case = int(sys.stdin.readline())
for _ in range(num_case):
    s = sys.stdin.readline().split()
    N = len(s) // 2
    pre_placed_rows = []
    pre_placed_cols = []
    n = 14
    for i in range(N):
        pre_placed_rows.append(int(s[2*i]))
        pre_placed_cols.append(int(s[2*i+1]))  
    
    print(len(nQueens(n, pre_placed_rows, pre_placed_cols, N)))

Answer 1

Here's my take on it.

Translate columns 1-14 to 0-13 for internal use and translate back on output.

state keeps track of where the queens are. state[0] is the row of the queen in column 0, etc. A None indicates the column is empty.

nQueens() searches state for an empty column and then checks each row in that column for a valid move. If a move is valid, state is updated and nQueens() is called with the new state. After the recursive call is made, revert the change in state. That avoids the need to make a copy of the state.

The code yields the solutions rather than keeping a list.

printboard() is just to help with visualizing what is going on. Uncomment some of the calls if you're curious...but only do one of the smaller test cases because there is a lot of output.

I think your imvalid() might be off?

Put the driver code in main().

SIZE = 14


def nQueens(state):
    #printboard(state)
    
    # find an empty column (row==None)
    for col, row in enumerate(state):
        if row is None:
            break
         
    # all columns filled. so it is a solution
    else:
        yield state
            
    for row in range(SIZE):
        # check if row taken
        if not valid(state, row, col):
            continue

        # make the move
        #print(f"col[{col+1}] = {row+1}")
        state[col] = row
        yield from nQueens(state)
        
        # revert the move to continue the search
        state[col] = None
            

def valid(state, row, col):
    # is there a queen in this row
    if row in state: 
        return False

    # is there a queen on a diagonal
    for c,r in enumerate(state):
        if r is None:
            continue
            
        if col-row == c-r or col+row == c+r:
            return False
                
    return True


def printboard(state):
    print('  '+' '.join(str(i%10) for i in range(1, SIZE+1)))
    for r in range(SIZE):
        print((r+1)%10, end=' ')
        for c in range(SIZE):
            if state[c] == r:
                print('Q', end=' ')
            else:
                print('.', end=' ')
        print()

        
def main(data):
    num_cases = int(data.readline())

    for case_num in range(num_cases):
        # initialize state to all Nones.
        # None means the corresponding column is empty
        state = [None]*SIZE

        # fill in the preset queens
        preset = data.readline().split()
        while preset:
            row = int(preset.pop(0)) - 1
            col = int(preset.pop(0)) - 1
            state[col] = row

        #printboard(state)

        # count the number of solutions
        for n,sol in enumerate(nQueens(state)):
            #printboard(sol)
            pass

        print(n+1)
        

Test it like so:

test = """
3
1 1 2 9 3 6 4 10
1 9 12 8 8 5 2 7
2 10 10 3
""".strip()

main(io.StringIO(test))

Or call it like:

main(sys.stdin)