5
\$\begingroup\$

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)))
\$\endgroup\$
6
  • 5
    \$\begingroup\$ Welcome to the Code Review site. Please add a link to the programming challenge. \$\endgroup\$ – pacmaninbw Nov 1 '20 at 14:56
  • 1
    \$\begingroup\$ Hi pacman, I do not have the link, but I have attached the screenshot. The time allowed is 6000ms, 256MB memory. My current algo has taken more than a minute, but less than 10MB memory \$\endgroup\$ – Michael W Nov 2 '20 at 13:34
  • \$\begingroup\$ How much more than a minute did it take? And did it produce the correct results? \$\endgroup\$ – superb rain Nov 2 '20 at 18:34
  • \$\begingroup\$ Hi superb, it is taking 20minutes for the sample input. It produced the correct results \$\endgroup\$ – Michael W Nov 3 '20 at 5:06
  • 1
    \$\begingroup\$ It looks like nQueens() constructs all possible solutions but drops those that don't have a queen in the pre-placed squares. That's a lot of extra work. state should start with the pre-placed queens already filed in and nQueens() should try to fill in the rest. \$\endgroup\$ – RootTwo Nov 4 '20 at 0:26
4
\$\begingroup\$

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)
\$\endgroup\$
2
  • \$\begingroup\$ Thanks RootTwo. Can I know which line will iterate all of the columns? \$\endgroup\$ – Michael W Nov 11 '20 at 6:18
  • \$\begingroup\$ The columns are iterated using recursion. state is a list of the position (row) of the queen in each column. That is, state[0] is the row of the queen in column 0, state[1] is the row of the queen in column 1, etc. The for col, row in enumerate(state): loop searches state to find an empty column. Just before each recursive call to nQueens(), that empty column of state is filled in. When state is full, that means every column has a queen. That should happen at the 10th, 11th, or 12th level of recursion depending on the number of preset queens. \$\endgroup\$ – RootTwo Nov 11 '20 at 20:51

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.