CF 102697107 - The Man Machine

We have an (N times N) chessboard and want to place exactly (N) queens so that no two queens can attack each other. A queen attacks horizontally, vertically, and diagonally, so every pair of queens must occupy different rows, different columns, and different diagonals.

codeforcescompetitive-programming

CF 102697107 - The Man Machine

Rating: -
Tags: -
Solve time: 1m 29s
Verified: yes

Solution

Problem Understanding

We have an (N \times N) chessboard and want to place exactly (N) queens so that no two queens can attack each other. A queen attacks horizontally, vertically, and diagonally, so every pair of queens must occupy different rows, different columns, and different diagonals. The output is the total number of distinct placements satisfying these conditions. The input contains only (N).

The crucial constraint is (N \le 9). That is small enough for exhaustive search with aggressive pruning, but large enough that blindly trying every board configuration is far too expensive. If we independently choose a column for the queen in every row, there are (N^N) possible placements. At (N=9), that is (9^9 = 387{,}420{,}489) candidates. Checking all pairs of queens at every complete candidate would require up to (\binom{9}{2}=36) pair checks, giving roughly (13.95) billion pair checks. A five-second solution cannot afford that amount of work.

There are several small cases that expose common implementation mistakes. For (N=1), the only square contains the only queen, so the answer is 1.

1

The correct output is:

1

A careless implementation that only increments the answer after making another recursive call can accidentally return 0 here because there is no next row to process.

For (N=2), there is no valid arrangement.

2

The correct output is:

0

The two queens must use different rows and columns, but in either possible diagonal arrangement they attack each other. Code that checks only rows and columns would incorrectly count two arrangements.

For (N=3), the answer is also zero.

3

The correct output is:

0

This catches implementations that handle columns correctly but have an incorrect diagonal formula, especially when the two diagonal directions are confused.

The sample (N=8) has 92 valid arrangements.

8

The correct output is:

92

The official problem uses (N=9) as its largest board size and explicitly requires an actual algorithm rather than a hard-coded lookup table.

Approaches

The most direct brute force is to process the board row by row. For every row, try all (N) columns, continue until all (N) rows have queens, and then check whether the resulting board is valid. This is correct because every possible placement has exactly one choice of column for each row, so the search eventually considers every possible board.

The problem is the number of candidates. With (N) choices in each of (N) rows, the search contains (N^N) complete placements. At the maximum (N=9), that is 387,420,489 placements. Even though many of them are immediately invalid, a brute-force implementation that waits until a complete board before checking conflicts still pays for all of them.

We can improve the search immediately by observing that two queens can never share a row. Instead of considering arbitrary board positions, place exactly one queen in each row. We can also reject a partial placement as soon as the new queen attacks an earlier queen. This changes the search from all (N^N) row-column assignments to a pruned search over partial placements.

There is another useful observation. Once we process rows from top to bottom, the only information needed to decide where the next queen can go is which columns and diagonals are already occupied. We do not need the complete board. Since (N) is at most 9, each of these sets fits naturally inside an integer bitmask.

For each recursive state, one bit represents each column. A set bit in cols means that column already contains a queen. The other two masks represent the diagonals currently attacked by queens above the next row. We can calculate every legal column at once with bit operations.

Suppose full contains the lowest (N) bits set to 1. Then

available = full & ~(cols | diag1 | diag2)

contains exactly the columns where the next queen can be placed. Picking one set bit gives one legal choice. After placing the queen, the diagonal masks are shifted because the next row is one row lower.

The brute-force works because it examines every possible placement, but fails because it explores enormous numbers of placements that already violate a constraint. The observation that every invalid partial placement can be discarded immediately lets us search only the feasible prefixes. Bitmasks then make finding and updating the feasible columns constant-time.

Approach Time Complexity Space Complexity Verdict
Brute Force (O(N^N \cdot N^2)) (O(N)) Too slow
Backtracking with bitmasks (O(N!)) worst case (O(N)) Accepted

The (O(N!)) bound is a conservative upper bound for the pruned search because once rows are fixed, no column can be reused, leaving at most (N, N-1, \ldots, 1) column choices along any unpruned permutation. Diagonal pruning makes the actual search considerably smaller.

Algorithm Walkthrough

  1. Read (N) and create full = (1 << N) - 1. The lowest (N) bits are now 1, representing all board columns.
  2. Start a recursive search at row 0 with empty column and diagonal masks. At row (r), the recursive state represents a valid placement of queens in rows (0) through (r-1).
  3. Calculate available = full & ~(cols | diag1 | diag2). Every set bit in this mask represents a column that is not occupied and is not attacked diagonally by an earlier queen.
  4. If available is zero before all rows have been filled, stop this branch. There is no legal position for the next queen, so continuing cannot produce a solution.
  5. Take one available bit with bit = available & -available. This isolates the lowest legal column. The ordering is irrelevant because every legal branch must eventually be counted, and choosing one bit at a time simply enumerates them without duplication.
  6. Remove that bit from available, then recursively place a queen in that column. The column mask becomes cols | bit.
  7. Move the diagonal attacks to the next row. The first diagonal mask is shifted left and the second is shifted right:
next_diag1 = ((diag1 | bit) << 1) & full
next_diag2 = (diag2 | bit) >> 1

The shifts are necessary because a diagonal moves one column horizontally whenever we move one row downward.

  1. When row == N, every row contains exactly one queen and all three masks have guaranteed that no two queens conflict. Increment the answer by one.

Why it works

The invariant is that every recursive call represents a placement of queens in all previously processed rows with no two queens attacking each other. cols records every occupied column, while the two diagonal masks record every column attacked diagonally in the next row. Consequently, available contains exactly the legal positions for the next queen. Every valid complete placement chooses one of these legal positions at every row, so it is never pruned. Every invalid placement has two queens sharing a column or diagonal, and the conflicting position is removed from available before it can be chosen. Thus every valid board is counted exactly once, and no invalid board is counted.

Python Solution

import sys
input = sys.stdin.readline

def solve():
    n = int(input())

    full = (1 << n) - 1
    answer = 0

    def backtrack(row, cols, diag1, diag2):
        nonlocal answer

        if row == n:
            answer += 1
            return

        available = full & ~(cols | diag1 | diag2)

        while available:
            bit = available & -available
            available -= bit

            backtrack(
                row + 1,
                cols | bit,
                ((diag1 | bit) << 1) & full,
                (diag2 | bit) >> 1
            )

    backtrack(0, 0, 0, 0)
    print(answer)

if __name__ == "__main__":
    solve()

The full mask contains exactly the columns belonging to the board. For (N=4), it is 1111 in binary. For (N=8), it is 11111111. Masking with full after shifting diag1 prevents bits that moved beyond the right edge from affecting later rows.

The recursive function receives the current row and three masks. There is no separate board array because the board itself is unnecessary once the occupied columns and diagonals are known.

The expression available & -available isolates one set bit. This lets the loop process one legal column at a time. Subtracting that bit removes it from the remaining candidates without modifying the masks belonging to the current recursive state.

The left shift and right shift in the recursive call are easy places to make an off-by-one mistake. A queen placed in column (c) attacks columns (c-1) and (c+1) on the next row, so the two diagonal directions must move in opposite directions. The & full after the left shift discards attacks that leave the board.

The base case checks row == n, rather than row == n - 1. At that point all (N) rows have received queens, so exactly one complete solution has been constructed. This also handles (N=1) correctly.

Python integers do not overflow, so the answer can be accumulated directly. For this problem the largest answer is only 352 at (N=9), but using an integer counter is still the natural implementation.

Worked Examples

Sample 1

For (N=8), one valid placement is represented by the row-to-column sequence 15863724. The following table traces that particular successful branch using zero-based bit positions. The full search explores many other branches as well, including branches that are pruned before reaching the eighth row.

Row Available Chosen bit Columns after choice Diag1 after choice Diag2 after choice
0 11111111 00000001 00000001 00000010 00000000
1 11111100 00010000 00010001 00100100 00001000
2 11000010 10000000 10010001 01001000 01000100
3 00100010 00100000 10110001 11010000 00110010
4 00001100 00000100 10110101 10101000 00011011
5 01000000 01000000 11110101 11010000 00101101
6 00000010 00000010 11110111 10100100 00010111
7 00001000 00001000 11111111 01011000 00001111

The chosen columns correspond to 1, 5, 8, 6, 3, 7, 2, 4 when converted back to one-based positions. At every row, the chosen bit is contained in available, so it conflicts with neither an occupied column nor either diagonal. The complete search finds 92 such leaves, giving the sample output of 92. The known count for the eight-queens problem is 92.

Sample 2

For (N=1), there is only one column and one row.

Row Available Chosen bit Result
0 1 1 Complete placement

After choosing the only bit, the recursion enters row == n, so the answer increases from 0 to 1. No diagonal conflict is possible because there is only one queen. This demonstrates why the base case must count a completed board immediately rather than expecting another row.

Complexity Analysis

Measure Complexity Explanation
Time (O(N!)) worst case At most (N!) row-by-row column permutations are considered, with constant-time mask operations at each state
Space (O(N)) The recursion depth is (N), and each level stores only a constant number of integers

With (N \le 9), the factorial upper bound is small enough for the 5-second limit, while the diagonal pruning makes the actual number of recursive states much smaller. The solution also uses only a few integers per recursion level, so it is far below the 256 MB memory limit.

Test Cases

The problem accepts only one (N) per execution, so the test helper below runs each input independently. The "all-equal values" category does not apply because the input contains a single scalar rather than an array.

import sys
import io

def count_queens(n):
    full = (1 << n) - 1
    answer = 0

    def backtrack(row, cols, diag1, diag2):
        nonlocal answer

        if row == n:
            answer += 1
            return

        available = full & ~(cols | diag1 | diag2)

        while available:
            bit = available & -available
            available -= bit

            backtrack(
                row + 1,
                cols | bit,
                ((diag1 | bit) << 1) & full,
                (diag2 | bit) >> 1
            )

    backtrack(0, 0, 0, 0)
    return answer

def solve():
    n = int(input())
    print(count_queens(n))

def run(inp: str) -> str:
    old_stdin = sys.stdin
    old_stdout = sys.stdout

    sys.stdin = io.StringIO(inp)
    sys.stdout = io.StringIO()

    try:
        solve()
        return sys.stdout.getvalue()
    finally:
        sys.stdin = old_stdin
        sys.stdout = old_stdout

# Provided sample 1
assert run("8\n") == "92\n", "sample 1"

# Provided sample 2
assert run("1\n") == "1\n", "sample 2"

# Minimum size
assert run("1\n") == "1\n", "minimum-size board"

# Smallest impossible board
assert run("2\n") == "0\n", "two queens cannot be placed safely"

# Boundary and diagonal correctness
assert run("4\n") == "2\n", "four queens has exactly two solutions"

# Maximum allowed size
assert run("9\n") == "352\n", "maximum-size board"
Test input Expected output What it validates
1 1 Minimum size and correct base case
2 0 Smallest impossible board and diagonal conflicts
4 2 Basic non-trivial search and both diagonal directions
9 352 Maximum allowed input and pruning performance

Edge Cases

For (N=1), the exact input is 1. The initial state has available = 1, so the only queen is placed. The recursive call reaches row == 1, increments the answer, and returns 1. There is no special-case branch in the implementation, which is preferable because the normal recursion already handles the smallest board.

For (N=2), the exact input is 2. The first queen can be placed in either column. Suppose it is placed in the first column. On the next row, the first column is blocked by the column mask and the second column is blocked by a diagonal mask, leaving no available position. The symmetric branch fails in the same way, so the final count is 0. This catches the common mistake of tracking columns but forgetting diagonals.

For (N=3), the exact input is 3. Some branches survive for two rows, but every possible third-row position is eventually blocked. The answer is 0. In particular, the two diagonal directions must be tracked independently because a queen can attack the next row from either side.

For (N=4), the exact input is 4. The search finds exactly two arrangements, commonly represented by the row-to-column sequences 2413 and 3142. The bitmask search counts both because it explores every legal column at each row, while branches that violate a diagonal are discarded before another recursive level is created.

For (N=8), the exact input is 8, and the output is 92. This is large enough to exercise substantial backtracking while remaining comfortably inside the constraint. A solution that repeatedly scans the entire board to check conflicts still works at this size in some languages, but the bitmask representation removes that repeated scanning and scales cleanly to the required maximum.

For (N=9), the exact input is 9, and the output is 352. This is the boundary case that matters most for performance. The implementation does not allocate an (N \times N) board, does not copy partial boards during recursion, and checks all legal columns with a handful of integer operations. The known total for nine queens is 352.