Queens, Sudoku, List monad

search.hs

import Control.Monad
import Data.List.HT
import Data.List
import System.Random.Shuffle
import Control.Monad.Random.Class
import qualified Data.Set as S

bitStringsN1s :: Int -> Int -> [[Int]]
bitStringsN1s n maxLength = do
  bitString <- replicateM maxLength [0, 1]
  guard $ length (filter (== 1) bitString) <= n
  return bitString

bitStringsNTrues :: Int -> Int -> [[Bool]]
bitStringsNTrues n maxLength = do
  bitString <- replicateM maxLength [True, False]
  guard $ length (filter id bitString) <= n
  return bitString

bitString11 :: Int -> [[Bool]]
bitString11 maxLength = do
  bitString <- replicateM maxLength [True, False]
  guard $ (length . filter id) bitString == 1
  return bitString

noAttacks :: [[Bool]] -> Bool
noAttacks board =
    all verifyQueens rows
  && all verifyQueens columns
  && all verifyQueens diagonal1
  && all verifyQueens diagonal2
  where rows = transposed
        columns = board
        diagonal1 = shear board
        diagonal2 = shear . map reverse $ board
        transposed = transpose board

verifyQueens :: [Bool] -> Bool
verifyQueens xs = (length . filter id) xs <= 1

countQueens :: Int -> [[Bool]] -> Bool
countQueens n xss = (length . filter id . concat) xss == n

queens :: Int -> [[[Bool]]]
queens n = do
  board <- replicateM n (bitStringsNTrues 1 n)
  guard $ noAttacks board
  guard $ countQueens n board
  return board

pretty :: [[Bool]] -> String
pretty bss =
  unlines . map (\bs -> map (\b -> if b then 'Q' else '.') bs) $ bss

randomRow :: (MonadRandom m) => Int -> m [Int]
randomRow n = shuffleM [1..n]

randomMatrix :: (MonadRandom m) => Int -> m [[Int]]
randomMatrix n = replicateM n (randomRow n)

boxAt :: [[Int]] -> Int -> Int -> Int -> [Int]
boxAt board boxSize x y =
  let xs = [x..x + (boxSize - 1)]
      ys = [y..y + (boxSize - 1)]
      coords = [(x1, y1) | x1 <- xs, y1 <- ys]
  in map (\(x2, y2) -> board !! x2 !! y2) coords

columnAt :: [[Int]] -> Int -> [Int]
columnAt = (!!)

rowAt :: [[Int]] -> Int -> [Int]
rowAt box n = map (!! n) box

verify :: [Int] -> Bool
verify xs = S.fromDistinctAscList [1..n] == S.fromList xs
  where n = length xs

sudoku :: (MonadRandom m) => Int -> m [[Int]]
sudoku n = do
  board <- randomMatrix n
  let rows = map (rowAt board) [0..(n - 1)]
  let cols = map (columnAt board) [0..(n - 1)]
  let boxes = map (uncurry (boxAt board step)) boxIndices
  if all verify rows && all verify cols && all verify boxes
   then return board
   else sudoku n
  where step = (truncate . sqrt . fromIntegral) n
        boxIndices = map (\x -> (x, x)) [0, step .. n]