n-queens

require 'set'

def build_board(n, queens)
  board = Array.new(n) { "." * n }
  queens.each { |spot|
    board[spot / n][spot % n] = "Q"
  }
  return board
end

def search_n_queens(n, spaces, sets, queens, solutions)
  if (n == queens.length)
    solutions.push(queens)
    return
  end
  return if (spaces.size < n - queens.length)
  return if (sets.values.any? { |set| set.size >= n })

  sets['rows'].each { |row|
    spaces -= Set.new((row*n)...(row*n + n))
  }
  sets['cols'].each { |col|
    column_set = Set.new
    (0...n).each { |row| column_set.add(row*n + col) }
    spaces -= column_set
  }
  sets['forw'].each { |forw|
    forward_set = Set.new
    [0, forw - (n - 1)].max.upto([forw, n - 1].min) { |cur|
      forward_set.add(cur*n + (forw - cur))
    }
    spaces -= forward_set
  }
  sets['back'].each { |back|
    back_set = Set.new
    row = [0, (n - 1) - back].max
    diff = back - (n - 1)
    while (row < n and row + diff < n)
      back_set.add(row*n + (row + diff))
      row += 1
    end
    spaces -= back_set
  }

  until spaces.empty?
    space = spaces.first
    spaces -= [space]
    row = space / n
    col = space % n
    diff = row - col
    new_sets = {
      'rows' => sets['rows'] + [row],
      'cols' => sets['cols'] + [col],
      'forw' => sets['forw'] + [row + col],
      'back' => sets['back'] + [(n - 1) - (row - col)]
    }
    search_n_queens(n, spaces.clone, new_sets, queens + [space], solutions)
  end
end

def solve_n_queens(n)
  solutions = Array.new
  sets = {
    'rows' => Set.new,
    'cols' => Set.new,
    'forw' => Set.new,
    'back' => Set.new
  }
  spaces = Set.new(0...n**2)
  search_n_queens(n, spaces, sets, Array.new, solutions)
  return solutions.map {|solution| build_board(n, solution)}
end

solve_n_queens(8).each { |board|
  puts board
  puts ""
}