queens.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns, TypeApplications #-}
import Control.Monad
import Data.List
import Diagrams.Prelude
import Diagrams.Backend.SVG.CmdLine
import Text.Printf

type Board = [Int]

queens :: Int -> (Board, Int) -> [Board]
queens n (board, !count)
  | n == count = return board
  | otherwise = do
      candidate <- [0 .. n - 1]
      let yDiffs = map (abs . subtract candidate) board
          xDiffs = map (abs . subtract count) indices
      guard $
        candidate `notElem` board
        && not (or $ zipWith (==) xDiffs yDiffs)
      queens n (candidate : board, count + 1)
 where
  indices = [count - 1, count - 2 ..]


-- Render to text
renderBoard :: Int -> Board -> String
renderBoard n = unlines . transpose . map renderRow
 where
  renderRow x = [if i == x then 'Q' else '.' | i <- [0 .. n - 1]]


-- Render to diagram
renderBoard' :: Int -> Board -> Diagram B
renderBoard' n b =
  vcat $ map hcat $ transpose (map renderRow b)
 where
  lw' = 0.2
  renderRow x =
    [ if i == x then
        circle 0.3 # fc black # lw lw'
        <> square 1 # bg white # lw lw'
      else
        square 1 # bg white # lw lw'
    | i <- [0 .. n - 1]
    ]


main :: IO ()
main =
  mainWith solveAndRender
 where
  margin :: Double
  margin = 0.7
  solveAndRender :: Int -> Diagram B
  solveAndRender n =
    let solutions = queens n ([], 0)
        count = length solutions
        root = truncate @Double . sqrt $ fromIntegral count
        diagram =
          vsep margin
          . map (hsep margin)
          . chunk root
          . map (\b -> b # alignTL `atop` square (width b) # lw 2 # alignTL )
          $ map (renderBoard' n) solutions
    in vcat . map center $
       [ diagram
       , text (printf "Solutions: %d" count) # fontSizeN 0.1 <> square (fromMeasured 1 10 (normalized 0.9)) # lw 0
       ]
  chunk :: Int -> [a] -> [[a]]
  chunk _ [] = []
  chunk c xs = take c xs : chunk c (drop c xs)