168 lines
2.6 KiB
Go
168 lines
2.6 KiB
Go
package n_queens
|
|
|
|
import (
|
|
"bytes"
|
|
)
|
|
|
|
type board [][]bool
|
|
|
|
func newBoard(size int) board {
|
|
b := make([][]bool, size)
|
|
for i := range b {
|
|
b[i] = make([]bool, size)
|
|
}
|
|
|
|
return board(b)
|
|
}
|
|
|
|
func (b board) size() (int, int) {
|
|
return len(b[0]), len(b)
|
|
}
|
|
|
|
func (b board) get(x, y int) bool {
|
|
return b[y][x]
|
|
}
|
|
|
|
func (b board) has(x, y int) bool {
|
|
width, height := b.size()
|
|
return x >= 0 && y >= 0 && x < width && y < height
|
|
}
|
|
|
|
func (b board) set(x, y int, value bool) {
|
|
b[y][x] = value
|
|
}
|
|
|
|
type position struct{ x, y int }
|
|
|
|
type step struct {
|
|
position int
|
|
blocked []position
|
|
}
|
|
|
|
type state struct {
|
|
board board
|
|
steps []step
|
|
n int
|
|
}
|
|
|
|
func newState(n int) *state {
|
|
return &state{
|
|
board: newBoard(n),
|
|
n: n,
|
|
steps: make([]step, 0, n),
|
|
}
|
|
}
|
|
|
|
func (s *state) options() []int {
|
|
y := len(s.steps)
|
|
if y >= s.n {
|
|
return []int{}
|
|
}
|
|
|
|
results := []int{}
|
|
for x := range s.n {
|
|
if s.board.get(x, y) {
|
|
continue
|
|
}
|
|
|
|
results = append(results, x)
|
|
}
|
|
|
|
return results
|
|
}
|
|
|
|
func (s *state) forward(x int) {
|
|
blocked := []position{}
|
|
y := len(s.steps)
|
|
|
|
// Vertical
|
|
for i := range s.n {
|
|
if !s.board.get(x, i) {
|
|
s.board.set(x, i, true)
|
|
blocked = append(blocked, position{x, i})
|
|
}
|
|
}
|
|
|
|
// Horizontal
|
|
for i := range s.n {
|
|
if !s.board.get(i, y) {
|
|
s.board.set(i, y, true)
|
|
blocked = append(blocked, position{i, y})
|
|
}
|
|
}
|
|
|
|
// Forward vertical
|
|
for offset := range s.n {
|
|
p := position{offset, y + x - offset}
|
|
|
|
if s.board.has(p.x, p.y) && !s.board.get(p.x, p.y) {
|
|
s.board.set(p.x, p.y, true)
|
|
blocked = append(blocked, p)
|
|
}
|
|
}
|
|
|
|
// Backward vertical
|
|
for offset := range s.n {
|
|
p := position{offset, y - x + offset}
|
|
|
|
if s.board.has(p.x, p.y) && !s.board.get(p.x, p.y) {
|
|
s.board.set(p.x, p.y, true)
|
|
blocked = append(blocked, p)
|
|
}
|
|
}
|
|
|
|
s.steps = append(s.steps, step{
|
|
position: x,
|
|
blocked: blocked,
|
|
})
|
|
}
|
|
|
|
func (s *state) backward() {
|
|
top, rest := s.steps[len(s.steps)-1], s.steps[:len(s.steps)-1]
|
|
s.steps = rest
|
|
|
|
for _, p := range top.blocked {
|
|
s.board.set(p.x, p.y, false)
|
|
}
|
|
}
|
|
|
|
func (s *state) solution() ([]string, bool) {
|
|
if len(s.steps) != s.n {
|
|
return nil, false
|
|
}
|
|
|
|
blank := bytes.Repeat([]byte{'.'}, s.n)
|
|
|
|
solution := make([]string, s.n)
|
|
for i := range solution {
|
|
row := bytes.Clone(blank)
|
|
row[s.steps[i].position] = 'Q'
|
|
solution[i] = string(row)
|
|
}
|
|
|
|
return solution, true
|
|
}
|
|
|
|
func SolveNQueens(n int) [][]string {
|
|
s := newState(n)
|
|
solutions := [][]string{}
|
|
|
|
var dfs func()
|
|
dfs = func() {
|
|
if soln, ok := s.solution(); ok {
|
|
solutions = append(solutions, soln)
|
|
return
|
|
}
|
|
|
|
for _, opt := range s.options() {
|
|
s.forward(opt)
|
|
dfs()
|
|
s.backward()
|
|
}
|
|
}
|
|
|
|
dfs()
|
|
|
|
return solutions
|
|
}
|