Problem statement

https://binarysearch.com/problems/N-Queens-Puzzle/

Solution

It is variation of Leetcode 0051. N-Queens. Here we given partial fill, so use bitmasks c0, c1, c2, c3 for rows, columns and two types of diagonals.

Complexity

It is still potentially O(n!) for time and O(n) for space.

Code

class Solution:
    def solve(self, M):
        def dfs(c0, c1, c2, c3):
            if c0 + 1 == 1 << n: return True
            for i in range(n):
                if not c0 >> i & 1: break

            for j in range(n):
                if c0 & 1<<i or c1 & 1<<j or c2 & 1<<i-j+n or c3 & 1<<i+j: continue
                if dfs(c0 ^ 1<<i, c1 ^ 1<<j, c2 ^ 1<<i-j+n, c3 ^ 1<<i+j): return True
            return False

        n = len(M)
        c0, c1, c2, c3 = 0, 0, 0, 0
        for i, j in product(range(n), range(n)):
            if M[i][j] == 1:
                if c0 & 1<<i or c1 & 1<<j or c2 & 1<<i-j+n or c3 & 1<<i+j: return False
                c0 ^= 1<<i
                c1 ^= 1<<j
                c2 ^= 1<<i-j+n
                c3 ^= 1<<i+j
        
        return dfs(c0, c1, c2, c3)