Counting N-Queens Solutions

		
Click to reveal / hide.
		
fn countSolutions(board: []u8, c: usize, n: usize) usize {
    if (c == n) return 1;

    var count: usize = 0;
    for (0..n) |r| {
	if (isSafe(board, r, c)) {
	    board[c] = @intCast(r);
	    count += countSolutions(board, c+1, n);
	}
    }
    return count;
}
	    

		
Click to reveal / hide.
		
fn isSafe(board: []u8, rj: usize, cj: usize) bool {
    for (board[0..cj], [0..cj]) |ri, ci| {
	if (ri == rj or
	    @max(ri, rj) - @min(ri, rj) == cj - ci)
	    return false;
    }
    return true;
}
	    

    
Click to reveal / hide.
    
fn fill(n: usize, board: []usize) void {
    for (0..n) |i|
	board[i] = i;
}

fn uniform_monte_carlo(n: usize, T: usize) usize {
    var rng = std.Random.DefaultPrng.init(0);
    const rand = rng.random();

    var board: [n]usize = undefined;
    fill(n, &board);

    // Enables O(n) safe check.
    var diags: [4 * n - 2]usize = undefined; 

    var S: usize = 0;
    for (0..T) |t| {
        rand.shuffle(usize, &board);
        if (fastSafe(n, t, &board, &diags))
            S += 1;
    }

    return S;
}
    

    
Click to reveal / hide.
    
inline fn fastSafe(
    n: usize,
    t: usize,
    board: []usize,
    diags: []usize,
) bool {
    for (0..n) |c| {
        const r = board[c];
        const a = r + c;
        const b = 3 * n - 2 + r - c;

        if (diags[a] == t or diags[b] == t) {
            return false;
        }

        diags[a] = t;
        diags[b] = t;
    }
    return true;
}

While the estimates look reasonable and were computed very quickly, the confidence intervals for \(n\ge 15\) are so large that one is not really inclined to believe the estimate, especially as I was unable to compute the ground-truth number of solutions for \(n=16\) at all. To tighten the confidence intervals we need to increase \(T\), but how much?

Analysis

Say we were content with merely having high confidence in the estimated order of magnitude. That is, we would like \(T\) large enough so that \(\hat{p}\) is with at least 95% probability within a factor 10 of the true \(p\). We can deduce how many samples we need from the Chernoff bound \[\operatorname{Pr}\left[|\hat{p} - p| \geq \epsilon p\right] \leq 2\exp\left(-\frac{\epsilon^2 p T}{3}\right) \leq \delta,\] with \(\epsilon = 9\) and \(\delta=0.05\). This tells us that the probably-correct-order-of-magnitude estimate is achievable with \[T\geq \frac{3}{\epsilon^2 p}\log\left(\frac{2}{\delta}\right)\] samples. But the probability \(p\) appears in the denominator of the lower bound, which is precisely the thing we are trying to estimate anyway, and has no closed form outside of an asymptotic result by Simkin (2021): \[p(n)=\frac{\left((1+o(1))ne^{-\alpha})\right)^n}{n!}\approx \frac{(ne^{-\alpha})^n}{n!}\approx\frac{e^{-(\alpha-1)n}}{\sqrt{2\pi n}},\] with \(\alpha\approx 1.944\). Unfortunately, it follows that we need an exponential amount of samples to estimate the correct order of magnitude with uniform sampling.

Conclusion

The fundamental problem of integration via uniform random sampling for \(n\)-queens is that the proportion of solutions is so small that for very large \(n\) it becomes unlikely to even see a single solution at all. This is something importance sampling will address.

3. Annealed Importance Sampling

Let's recap. We wanted to find the number of ways of safely placing \(n\) queens on an \(n\)-by-\(n\) chessboard. Exhaustive counting is too slow, and so is uniform Monte Carlo integration of the indicator function of the solution set.

Motivation

There are a few important facts that will let us get around the problem of vanishing \(p(n)\). Here is fact number one, stated abstractly.

1. Given positive numbers \(A\lt B \lt C\), the fractions involved in the equation \[\frac{A}{C}=\frac{A}{B}\frac{B}{C}\] are larger on the right hand side than the left hand side.

Applying this fact to our problem, we can estimate \(Q(n)/n!\) by constructing a sequence of intermediate fractions whose size we may better estimate individually. To do that, we need a way to climb monotonically from \(Q(n)\) to \(n\).

Interpolating between \(Q(n)\) and \(n!\)

Recall that these \(Q(n)\) and \(n!\) may be understood as discrete integrals of indicator functions on \(\mathcal{S}\) and \(\mathcal{P}\), respectively. This hint leads us to fact number two.

2. An interpolation from \(\mathbf{1}_\mathcal{P}\) to \(\mathbf{1}_\mathcal{S}\) is provided by the Boltzmann function \[X \mapsto \exp{(-\beta f(X))}\] when the inverse temperature \(\beta\) varies from \(0\) to \(\infty\), and \(f\) is a positive energy function that is zero on the solution set \(\mathcal{S}\).[2]

Play with the \(\beta\) slider to see how it tempers the original function, singling out the two solutions (\(n=4\)).

Our choice of energy function \(f\) simply returns the number of diagonal threats. This energy function is nice for a reason that will make sense later.

As a side note, if there are \(k\) queens in a diagonal, then that diagonal contributes \(k(k-1)/2\) threats to the energy function, and the total number of threats (energy) is the sum over all diagonals.

inline fn pairs(comptime T: type, k: T) T {
    return @divExact(k * (k - 1), 2);
}

inline fn threats(diags: []isize) isize {
    var sum: isize = 0;
    for (diags) |d|
        sum += pairs(isize, d);
    return sum;
}

Returning to the point, let's discuss how to use the intermediate functions to build a ladder from \(Q(n)\) to \(n!\). The discrete integral of the Boltzmann function is a normalisation constant \[Z(\beta)=\sum_{X\in \mathcal{P}}\exp{(-\beta f(X))}\] that yields the desired interpolation between \(Q(n)\) and \(n!\) since \(Q(n)=Z(\infty)\lt \dots \lt Z(0)= n!\). To apply this fact in our estimation, we create a temperature schedule \(\beta_1 \lt \beta_2\lt \dots \lt \beta_K\approx \infty\), and compute the number of solutions as follows: \[\frac{Q(n)}{n!}=\prod_{k=1}^{K-1} \frac{Z(\beta_{k+1})}{Z(\beta_k)}:= \prod_{k=1}^{K-1} w_k.\]

Now, the schedule design matters. If \(\beta_k\) and \(\beta_{k+1}\) are close enough, then so are \(Z(\beta_{k})\) and \(Z(\beta_k)\), thus we benefit from balancing the factors. Too close, and we waste our time on factors with no impact.

Designing a well-balanced temperature schedule depends on how we estimate the fractions, so we turn our attention to this problem next.

Importance sampling

Instead of uniform sampling, we switch to Boltzmann sampling. To make the switch, introduce the Boltzmann distribution \[\pi_{\beta}(X)=\frac{\exp{(-\beta f(X))}} {Z(\beta)}\] in the formula as follows. \[w_k=\sum_{X\in\mathcal{P}}\exp{\left(-(\beta_{k+1} - \beta_k)f(X)\right)}\pi_{\beta_k}(X).\] This lets us estimate the very same ratio by drawing M independent samples \(X_1,\dots, X_M\) according to \(\pi_{\beta_k}\) and computing the estimator \[\hat{w}_k=\frac{1}{M}\sum_{m=1}^M \exp{(-(\beta_{k+1}-\beta_k)f(X_m))}.\]

This switch of sampling distributions an example of importance sampling. The new estimator is unbiased in the sense that \[\mathbb{E}_{\pi_{\beta_k}}\left[\hat{w}_{k}\right]=w_k,\] and it has lower variance than a uniform estimator since lower energy elements are sampled more often under \(\pi_\beta\) than under the uniform distribution. Either way, our final estimator of \(Q(n)/n!\) is \[\hat{p}(n)=\prod_{k=1}^{K-1}\hat{w}_{k}.\]

But doesn't this mean we have to incur the cost of computing the distribution \(\pi_\beta\) with its \(n!\)-sized support? Actually, the answer is no.

Metropolis algorithm

To sample according to \(\pi_\beta\), we set up a Markov chain whose limiting distribution is theoretically \(\pi_\beta\). Press the 🔥 button to make a random Markov chain transition: a column swap.

			
Click to reveal / hide.
			
// convenient casting
inline fn F64(x: anytype) f64 {
    return @floatFromInt(x);
}

// annealed importance sampling
fn ais(
    comptime n: usize,
    comptime M: usize,
) f64 {

    // set up random number generator
    var rng = std.Random.DefaultPrng.init(0);
    const rand = rng.random();

    // set up board.
    var board: [n]usize = undefined;
    fill(n, &board);
    var diags: [4 * n - 2]isize = undefined;
    fillDiags(n, &board, &diags);

    // initialize schedule and estimate
    var beta: f64 = 0;
    var estimate: f64 = 0;

    while (true) {

        var energies: [M]usize = undefined;
        metropolis(n, M, rand, &board, &diags, beta, &energies);

        const delta_beta = compute_delta_beta(&energies);
        beta += delta_beta;

        var w: f64 = 0;
        for (energies) |f|
            w += @exp(-delta_beta * F64(f));

        estimate += @log(w) - @log(F64(M));

        var sum: usize = 0;
        for (energies) |f| {
            sum += if (f == 0) 1 else 0;
        }
        if (F64(sum) >= 0.999 * F64(M)) break;
    }

    for (1..(n + 1)) |i|
        estimate += @log(F64(i));

    return @exp(estimate);
}
		    

		    
Click to reveal / hide
		    
inline fn metropolis(
    comptime n: usize,
    comptime M: usize,
    rand: std.Random,
    board: []usize,
    diags: []isize,
    beta: f64,
    energies: []usize,
) void {
    var energy: usize = @intCast(threats(diags));
    for (0..M) |m| {
        energies[m] = energy;

        // Pick two columns to swap
        const c1 = rand.uintAtMost(usize, n - 1);
        const c2 = (c1 + 1 + rand.uintAtMost(usize, n - 2)) % n;
        const r1 = board[c1];
        const r2 = board[c2];

        // Note indices of occupied diagonals.
        const offset = 3 * n - 2;
	const diag_idxs: [8]usize = .{
	    r1 + c1,
	    offset + r1 - c1,
	    r2 + c2,
	    offset + r2 - c2,
	    r2 + c1,
	    offset + r2 - c1,
	    r1 + c2,
	    offset + r1 - c2,
	};

	// Change in number of pieces on diagonal due to swap.
        const piece_deltas: [8]isize = .{ -1, -1, -1, -1, 1, 1, 1, 1 };

        // Compute energy change due to swap
        const energy_delta = step(diags, diag_idxs, piece_deltas);

        // If swap accepted, commit and update board, energy, diagonals 
        if (@log(rand.float(f64)) < @min(0, -beta * F64(energy_delta))) {
            const tmp = board[c1];
            board[c1] = board[c2];
            board[c2] = tmp;

            energy = @intCast(@as(isize, @intCast(energy)) + energy_delta);

            for (0..8) |i|
                diags[diag_idxs[i]] += piece_deltas[i];
        }
    }
}

			
Click to reveal / hide.
inline fn step(
    diags: []isize,		// how many queens per diagonal
    diag_idxs: [8]usize,	// diagonals involved in the swap
    piece_deltas: [8]isize,	// whether those diagonals lose or gain a piece
) isize {
    var energy_delta: isize = 0;

    // Count the total gain or loss of pieces per diagonal
    var diag_read: [8]bool = .{false} ** 8;
    for (0..8) |i| {
        if (diag_read[i]) continue;

        const idx = diag_idxs[i];
        var total_piece_delta = piece_deltas[i];

        // Look ahead to see if diagonal is involved with other piece
        for (i + 1..8) |j| {
            if (diag_idxs[j] == idx and !diag_read[j]) {
                total_piece_delta += piece_deltas[j];
                diag_read[j] = true;
            }
        }

        // Continue if number of pieces stays the same.
        if (total_piece_delta == 0) continue;

	// Compute change in number of threats
        const d_old = diags[idx];
        const d_new = d_old + total_piece_delta;
        energy_delta += pairs(isize, d_new) - pairs(isize, d_old);
    }

    return energy_delta;
}