leetcode-go/hard/52.go

package main

func totalNQueens(n int) (total int) {
	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 {
				total++
			}
			y-- // Back tracking
			x = putX[y]
			colLaid[x] = false
			mainLaid[y-x+n] = false
			viceLaid[x+y] = false
			x++
		}
	}
}

func main() {
	println(totalNQueens(4))
	println(totalNQueens(5))
	println(totalNQueens(6))
	println(totalNQueens(7))
}