Fast algorithm for n-rooks completion puzzle

Question

I am motivated by the post, Complexity of n-queens-completion. I am interested in completion problem of non-attacking rooks on a chessboard.

Input: Given a chessboard of size 𝑛×𝑛 with 𝑛−𝑘 rooks placed in non-attacking positions, and certain chessboard diagonals represented by multi-set of integers 𝐷 = { 𝑗₁, 𝑗₂, ..., 𝑗_𝑘 }

Output: Can we place additional 𝑘 rooks such all 𝑛 rooks are non-attacking and each new rook is placed on exactly one diagonal in 𝐷?

I suspect that the problem is NP-complete. Is there an efficient (fast) algorithm?

P.S. From comments, the multi-set of diagonals D defines parallel diagonals. For diagonal j, the sum of x and y coordinates of any cell on that diagonal is equal to j. This implies that there are no crossing diagonals. So, Only parallel diagonals (only diagonals, no anti-diagonals).

Answer 1

This is a constraint satisfaction problem, and more particularly a SAT problem, where each boolean variable corresponds to a square on the chess board, indicating whether it gets a rook or not.

The constraints then express that every non-occupied row must have exactly one rook, similarly every non-occupied column must have exactly one rook. The given diagonals must have the number of rooks specified for them. The constraints only need to reference squares that lie on at least one given diagonal and are not on a row/col that already has a rook. The other squares can be ignored, as it is clear they wont get a rook.

There are several packages for python that solve such kinds of problems, like scipy, numpy, and others. Here I will use the python-constraint package.

from constraint import Problem, ExactSumConstraint
from collections import Counter

# rooks = list of coordinates of pre-placed rooks
# diagonals = list of diagonals identified by row+col
#             It can have duplicates. Each given diagonal will
#             get exactly one new rook (not counting pre-placed rooks)
# Note that the size of the chess board is implied by the above input
# Returns a list of all rooks, including the pre-placed rooks, None if no solution

def solve(rooks, diagonals):
    n = len(rooks) + len(diagonals)
    rows = set(range(n)).difference(row for row, col in rooks)
    cols = set(range(n)).difference(col for row, col in rooks)
    diagcounter = Counter(diagonals)

    # cells that might receive a rook:
    #   ...are on rows and cols that don't have a rook yet, and
    #   ...are on one of the given diagonals
    cells = {
        (row, col) 
        for col in cols
        for row in rows
        if row + col in diagcounter
    }
    rowconstraints = [
        ([
            (row, col) for col in cols
            if (row, col) in cells
        ], 1)
        for row in rows
    ]
    colconstraints = [
        ([
            (row, col) for row in rows
            if (row, col) in cells
        ], 1)
        for col in cols
    ]
    diagconstraints = [
        ([
            (row, diag - row) for row in rows
            if diag - row in cols 
        ], numRooks)
        for diag, numRooks in diagcounter.items()
    ]
    constraints = rowconstraints + colconstraints + diagconstraints
    # Translate to how python-constraint package works 
    problem = Problem()
    problem.addVariables(cells, [0, 1])
    for cells, numRooks in constraints:
        problem.addConstraint(ExactSumConstraint(numRooks), cells)
    solution = next(problem.getSolutionIter(), None)
    if solution:
        return sorted([cell for cell, hasrook in solution.items() if hasrook] + rooks)

Here is an example run:

# Example input (see image below)
rooks = [(1, 2)]                    # row, col coordinates of each known rook
diagonals = [3, 7, 7, 7, 9, 10, 10] # diags identified by row+col sum

solution = solve(rooks, diagonals)
print(solution)

The input represents this chess board and constraints:

There is one pre-placed rook (indicated with R). The diagonals in 𝐷 are colored, and their frequency indicated with a number. For example, the yellow one needs to get 3 rooks and is identified by the number 7 since its squares have coordinates that add up to 7.

A solution to this test case is this:

... and the output of the above example run is:

[(0, 3), (1, 2), (2, 5), (3, 7), (4, 6), (5, 4), (6, 1), (7, 0)]