N-Queens II using backtracking is slow

The restricted set seems wasteful, both time- and space-wise. At the end of the successful recursion, n levels deep it grows to n^2 size, which drives the total complexity to O(n^3). And it is not really needed. It is much easier to check the availability of the square by looking at the queens already placed (please forgive the chess lingo; file stand for vertical, and rank for horizontal):

def square_is_safe(file, rank, queens_placed):
    for queen_rank, queen_file in enumerate(queens_placed):
        if queen_file == file:                      # vertical attack
            return false
        if queen_file - file == queen_rank - rank:  # diagonal attack
            return false
        if queen_file - file == rank - queen_rank:  # anti-diagonal attack
            return false
    return true

to be used in

def place_queen_at_rank(queens_placed, rank):
    if rank == n:
        total_count += 1
        return

    for file in range(0, n):
        if square_is_safe(file, rank, queens_placed):
            queens_placed.append(file)
            place_queen_at_rank(queens_placed, rank + 1)

    queens_placed.pop()

And there is a plenty of room for the optimization. For example, you may want to special-case the first rank: due to a symmetry, you only need to inspect a half of it (cutting execution time by the factor of 2).

For n ≤ 9 (the bound in the linked puzzle), it's enough to enumerate all valid positions for rooks and verify that there are no attacking diagonal moves.

import itertools


def is_valid(ranks):
    return not any(
        abs(f1 - f2) == abs(r1 - r2)
        for f1, r1 in enumerate(ranks)
        for f2, r2 in enumerate(ranks[:f1])
    )


def count_valid(n):
    return sum(map(is_valid, itertools.permutations(range(n))))


print(*(count_valid(i) for i in range(1, 10)), sep=",")

In this kind of problem, you have to focus on the algorithm first, not on the code.

In the following, I will focus on the algorithm, just giving an example in C++ to illustrate it.

One main issue is to be able to detect fast if a given position is already controlled or not by an existing Queen.

One simple possibility is to index the diagonals (for 0 to 2N-1), and to keep track in a array if the corresponding diagonals, antidiagonals or the columns are already controlled. Any way to index the diagonals or the antidiagonals will do the jobs. For a given (row, column) point, I use:

diagonal index = row + column
antidiagonal index = n-1 + col - row

In addition, I use a simple symmetry: it is only necessary to calculate the number of possibilities for a row index from 0 to n/2-1 (or n/2 if n is odd).

It is certainy possible to speed it up a little bit, by using other symmmetries. However, as it is, it looks quite fast enough for n values less or equal to 9.

Result:

2 : 0 time : 0.001 ms
3 : 0 time : 0.001 ms
4 : 2 time : 0.001 ms
5 : 10 time : 0.002 ms
6 : 4 time : 0.004 ms
7 : 40 time : 0.015 ms
8 : 92 time : 0.05 ms
9 : 352 time : 0.241 ms
10 : 724 time : 0.988 ms
11 : 2680 time : 5.55 ms
12 : 14200 time : 31.397 ms
13 : 73712 time : 188.12 ms
14 : 365596 time : 1046.43 ms

Here is the code in C++. As the code is quite simple, you should easily be able to convert it in Python.

#include <iostream>
#include <chrono>

constexpr int N_MAX = 14;
constexpr int N_DIAG = 2*N_MAX + 1;

class Solution {
public:
    int n;
    int Col[N_MAX] = {0};
    int Diag[N_DIAG] = {0};
    int AntiDiag[N_DIAG] = {0};
    
    int totalNQueens(int n1) {
        n = n1;
        if (n <= 1) return n;
        int count = 0;
        for (int col = 0; col < n/2; ++col) {
            count += sum_from (0, col);
        }
        count *= 2;
        if (n%2) count += sum_from (0, n/2);
        return count;
    }
    
    int sum_from (int row, int col) {
        if (Col[col]) return 0;
        int diag = row + col;
        if (Diag[diag]) return 0;
        int antidiag = n-1 + col - row;
        if(AntiDiag[antidiag]) return 0;
        if (row == n-1) return 1;
        int count = 0;
        Col[col] = 1;
        Diag[diag] = 1;
        AntiDiag[antidiag] = 1;
        for (int k = 0; k < n; ++k) {
            count += sum_from (row+1, k);
        }
        Col[col] = 0;
        Diag[diag] = 0;
        AntiDiag[antidiag] = 0;
        return count;
    }
};


int main () {
    int n = 1;
    while (n++ < N_MAX) {
        auto start = std::chrono::high_resolution_clock::now();
        Solution Sol;
        std::cout << n << " : " << Sol.totalNQueens (n) << " time : ";
        auto diff = std::chrono::high_resolution_clock::now() - start;
        auto duration = std::chrono::duration_cast<std::chrono::microseconds>(diff).count();
        std::cout << double(duration)/1000 << " ms" << std::endl;
    }
    return 0;
}

You can avoid checking for horizontal conflicts by only placing one queen per row. This also allows you to reduce the size of the diagonal conflict matrix by only flagging subsequent rows. Using a simple boolean flag list for column conflicts is also a time saver (as opposed to flagging multiple entries in a matrix)

here's an example as a generator of solutions:

def genNQueens(size=8):
    # setup queen coverage from each position {position:set of positions}
    reach = { (r,c):[] for r in range(size) for c in range(0,size) }
    for R in range(size):
        for C in range(size):
            for h in (1,-1): # diagonals on next rows
                reach[R,C].extend((R+i,C+h*i) for i in range(1,size))
            reach[R,C] = [P for P in reach[R,C] if P in reach]
    reach.update({(r,-1):[] for r in range(size)}) # for unplaced rows

    # place 1 queen on each row, with backtracking
    cols     = [-1]*size            # column of each queen (start unplaced)
    usedCols = [False]*(size+1)     # column conflict detection
    usedDiag = [[0]*(size+1) for _ in range(size+1)] # for diagonal conflicts
    r        = 0
    while r >= 0:
        usedCols[cols[r]] = False
        for ur,uc in reach[r,cols[r]]: usedDiag[ur][uc] -= 1
        cols[r] = next((c for c in range(cols[r]+1,size)
                        if not usedCols[c] and not usedDiag[r][c]),-1)
        usedCols[cols[r]] = True
        for ur,uc in reach[r,cols[r]]: usedDiag[ur][uc] += 1
        r += 1 if cols[r]>=0 else -1   # progress or backtrack
        if r<size : continue           # continue until all rows placed
        yield [*enumerate(cols)]       # return result
        r -= 1                         # backtrack to find more
        

output:

from timeit import timeit
for n in range(3,13):
    t = timeit(lambda:sum(1 for _ in genNQueens(n)), number=1)
    c = sum(1 for _ in genNQueens(n))
    print(f"solutions for {n}x{n}:", c, "time:",f"{t:.4g}")    

solutions for 3x3: 0 time: 0.000108
solutions for 4x4: 2 time: 0.0002044
solutions for 5x5: 10 time: 0.0004365
solutions for 6x6: 4 time: 0.0008741
solutions for 7x7: 40 time: 0.003386
solutions for 8x8: 92 time: 0.009881
solutions for 9x9: 352 time: 0.03402
solutions for 10x10: 724 time: 0.1228
solutions for 11x11: 2680 time: 0.5707
solutions for 12x12: 14200 time: 2.77
    

Just one change (remove r loop in gen) can make your solution AC.

The main reason is that your gen has argument row, and it will call itself using row + 1, so there is no need to iterate using for r in range(row, n):. It's unnecessary. Just Removing it, your solution is quite acceptable.(we need to need add else before nested call)

Following is the result: Before change:

1 1      1.8358230590820312e-05
2 0      5.7697296142578125e-05
3 0      0.00036835670471191406
4 2      0.0021448135375976562
5 10     0.02212214469909668
6 4      0.23602914810180664
7 40     3.0731561183929443

After change:

1 1      1.6450881958007812e-05
2 0      3.1948089599609375e-05
3 0      0.0001366138458251953
4 2      0.0002281665802001953
5 10     0.0008234977722167969
6 4      0.0028502941131591797
7 40     0.01242375373840332
8 92     0.05443763732910156
9 352    0.2279810905456543

It just uses 0.4% time of original version for the n = 7 case, and n = 8 can absolutely work.

class Solution:
    def totalNQueens(self, n: int) -> int:
        def genRestricted(restricted, r, c):
            restricted = set(restricted)
            for row in range(n): restricted.add((row, c))
            for col in range(n): restricted.add((r, col))
            movements = [[-1, -1], [-1, 1], [1, -1], [1, 1]]
            for movement in movements:
                row, col = r, c
                while 0 <= row < n and 0 <= col < n:
                    restricted.add((row, col))
                    row += movement[0]
                    col += movement[1]
            return restricted

        def gen(row, col, curCount, restricted):
            count, total_count = curCount, 0

            for c in range(col, n):
                if (row, c) not in restricted:
                    count += 1
                    if count == n: total_count += 1
                    else: total_count += gen(row + 1, 0, count, genRestricted(restricted, row, c))
                    count -= 1

            return total_count

        return gen(0, 0, 0, set())
if __name__ == '__main__':
    import time
    s = Solution()
    for i in range(1, 8):
        t0 = time.time()
        print(i, s.totalNQueens(i), '\t', time.time() - t0)

Of course, there are other enhancements can be made. But this is the biggest one.

For instance, you updated and created a new restricted/forbidden points after adding each point. BTW, I don't agree @user58697 for restricted, it's necessary based on your solution, as you need to clone and update to get a new one to avoid restore it in the recursive call loop.

BTW, following is my solution, just for your reference:

class Solution:
    def solveNQueens_n(self, n): #: int) -> List[List[str]]:
        cols = [-1] * n # index means row index
        self.res = 0
        usedCols = set() # this and cols can avoid vertical and horizontal conflict

        def dfs(r): # current row to fill in
            def valid(c):
                for r0 in range(r):
                    # (r0, c0), (r1, c1) in the (back-)diagonal, |r1 - r0| = |c1 - c0|
                    if abs(c - cols[r0]) == abs(r - r0):
                        return False
                return True
            if r == n: # valid answer
                self.res += 1
                return
            for c in range(n):
                if c not in usedCols and valid(c):
                    usedCols.add(c)
                    cols[r] = c
                    dfs(r + 1)
                    usedCols.remove(c)
                    cols[r] = -1

        dfs(0)
        return self.res

Ok, one thing I missed was that each row must have a queen. Very important observation. gen method has to be modified like this:

    def gen(row, col, curCount, restricted):
        if row == n: return 0
        
        count, total_count = curCount, 0
        
        for c in range(col, n):
            if (row, c) not in restricted:
                if count + 1 == n: total_count += 1
                total_count += gen(row + 1, 0, count + 1, genRestricted(restricted, row, c))
                    
        return total_count

It beats only ~20% submissions, so it's not perfect at all. Far from it.