leetcode-go/hard/51.go

package main

import (
	"fmt"
	"strings"
)

func solveNQueens(n int) (solutions [][]string) {
	putX := make([]int, n)
	colLaid := make([]bool, n)
	mainLaid := make([]bool, 2*n) // Range of y-x is [1-n, n-1], so y-x+n -> [1, 2n-1]
	viceLaid := make([]bool, 2*n-1)
	for y, x := 0, 0; ; { // The ith queen
		for ; x < n; x++ {
			if !(colLaid[x] || mainLaid[y-x+n] || viceLaid[x+y]) {
				putX[y] = x
				colLaid[x] = true
				mainLaid[y-x+n] = true
				viceLaid[x+y] = true
				x = 0 // Next step
				y++
				break
			}
		}
		if x == n || y == n { // Not valid or finished
			if y == 0 { // If y is 0 and x is n, then all solutions are found
				return
			} else if y == n {
				solutions = append(solutions, make([]string, 0))
				l := len(solutions)
				var sb strings.Builder
				for i := 0; i < n; i++ {
					sb.Reset()
					for j := 0; j < n; j++ {
						if putX[i] == j {
							sb.WriteByte('Q')
						} else {
							sb.WriteByte('.')
						}
					}
					solutions[l-1] = append(solutions[l-1], sb.String())
				}
			}
			y-- // Back tracking
			x = putX[y]
			colLaid[x] = false
			mainLaid[y-x+n] = false
			viceLaid[x+y] = false
			x++
		}
	}
}

func main() {
	fmt.Println(solveNQueens(9))
}