bhakti-thakur
81 lines
2.6 KiB
Python
81 lines
2.6 KiB
Python
def solve_n_queens_with_first(n, first_pos, find_all=False):
|
|
"""
|
|
Solve n-Queens given a pre-placed queen at first_pos (r, c), 0-indexed.
|
|
Returns:
|
|
- if find_all=False: one solution as a list of rows where board[r][c] = 1
|
|
- if find_all=True: list of all such solutions
|
|
If no solution: returns None (or empty list when find_all=True).
|
|
"""
|
|
r0, c0 = first_pos
|
|
# validate first position
|
|
if not (0 <= r0 < n and 0 <= c0 < n):
|
|
raise ValueError("first_pos must be within board bounds (0..n-1).")
|
|
|
|
# Prepare trackers
|
|
cols = set([c0])
|
|
diag1 = set([r0 - c0]) # r - c
|
|
diag2 = set([r0 + c0]) # r + c
|
|
board = [-1] * n # board[r] = c or -1
|
|
board[r0] = c0
|
|
|
|
solutions = []
|
|
|
|
def place(row):
|
|
# If row already has the placed queen, skip to next row
|
|
if row == n:
|
|
# a solution found; build matrix representation
|
|
mat = [[0]*n for _ in range(n)]
|
|
for rr in range(n):
|
|
c = board[rr]
|
|
if c != -1:
|
|
mat[rr][c] = 1
|
|
if find_all:
|
|
solutions.append(mat)
|
|
return False # continue searching
|
|
else:
|
|
solutions.append(mat)
|
|
return True # stop search (found one)
|
|
if board[row] != -1:
|
|
# fixed queen row, just advance
|
|
return place(row+1)
|
|
|
|
for c in range(n):
|
|
if c in cols or (row - c) in diag1 or (row + c) in diag2:
|
|
continue
|
|
# choose
|
|
board[row] = c
|
|
cols.add(c); diag1.add(row - c); diag2.add(row + c)
|
|
stop = place(row + 1)
|
|
# unchoose
|
|
cols.remove(c); diag1.remove(row - c); diag2.remove(row + c)
|
|
board[row] = -1
|
|
if stop and not find_all:
|
|
return True
|
|
return False
|
|
|
|
# Start from row 0
|
|
place(0)
|
|
|
|
if find_all:
|
|
return solutions # possibly empty list
|
|
else:
|
|
return solutions[0] if solutions else None
|
|
|
|
|
|
# ---------- Example usage ----------
|
|
if __name__ == "__main__":
|
|
# Example 1: n=8, first queen at row0,col0 (top-left)
|
|
n = 8
|
|
first = (0, 0) # 0-indexed
|
|
sol = solve_n_queens_with_first(n, first, find_all=False)
|
|
if sol is None:
|
|
print("No solution found.")
|
|
else:
|
|
print("One solution matrix (1 = queen):")
|
|
for row in sol:
|
|
print(row)
|
|
|
|
# Example 2: get all solutions with first queen at (0, 3)
|
|
# all_sols = solve_n_queens_with_first(8, (0, 3), find_all=True)
|
|
# print(f"Found {len(all_sols)} solutions.")
|