N-Queens Problem in Haskell using Genetic Algorithms

GeneticAlgorithms.hs

{-# LANGUAGE NoImplicitPrelude #-}

-- | N-Queens Problem: Place n queens on a chess board so that neither of them is threatened
-- This program will find a solution using a genetic algorithm

module GeneticAlgorithms where

import ClassyPrelude
import Data.Vector.Unboxed (ifoldl',update)
import System.Random


main :: IO ()
main = do
  (level, fittest, generations) <- nQueens 8 1000 1000
  print $ "Fittest: " <> show fittest
  print $ "Level:   " <> show level
  print $ "Gens:    " <> show generations
  drawBoard fittest


-- Vector Index = column, Vector Cell = row ... e.g. [0 4 2 3 4 5 7 7]
type Board = UVector Int

-- | Solution of N-Queens Problem
nQueens :: Int                  -- ^ width and height of chess board
        -> Int                  -- ^ size of the population
        -> Int                  -- ^ maximum generations before GA stops
        -> IO (Int, Board, Int) -- ^ (fitness level, fittest board, generation of fittest) 
nQueens boardSize popSize maxGens = do
  boards         <- randBoards popSize
  (fittest,left) <- evolve maxGens boards
  return (fitness fittest, fittest, maxGens - left)
  where
    randField :: IO Int
    randField = randomRIO (0,boardSize-1)
    randBoards :: Int -> IO (Vector Board)
    randBoards n = replicateM n (replicateM boardSize randField)
    mate :: Board -> Board -> IO Board
    mate mom dad = (\f -> take f mom <> drop f dad) <$> randField
    evolve :: Int -> Vector Board -> IO (Board, Int)
    evolve maxGens pop
      | fitness fittest == 0 = return (fittest, maxGens)
      | maxGens < 1          = return (fittest, maxGens)
      | otherwise            = replicateM popSize fitChild >>= evolve (maxGens - 1)
      where
        sortedPop = sortWith fitness pop
        fittest = headEx sortedPop
        fitChild :: IO Board
        fitChild = do
          momIdx <- (popSize - 1 -) <$> randE (0, popSize - 1)
          dadIdx <- (popSize - 1 -) <$> randE (0, popSize - 1)
          child  <- mate (sortedPop!momIdx) (sortedPop!dadIdx)
          oneInThousand <- (1 ==) <$> randomRIO (1,1000 :: Int)
          if oneInThousand
          then mutate child
          else return child
        mutate :: Board -> IO Board
        mutate board = do
          i <- randField
          p <- randField
          return $ update board [(i,p)]

-- | Fitness of a board (number of threats)
-- c = column, r = row
fitness :: UVector Int -> Int
fitness board = ifoldl' (\acc c r -> acc + straight c r + diagonal c r) 0 board
  where
    after c = drop (c+1) board
    straight c r
      | any (== r) (after c) = 1
      | otherwise            = 0
    diagonal c r
      | ifoldl' (\found c' r' -> found || abs (r-r') == (c'+1)) False (after c) = 1
      | otherwise     = 0

-- | Random number in range calculated with e^x
randE :: (Int,Int) -> IO Int
randE (low,up) = do
  x <- randomRIO (0,4) :: IO Double
  return $ round (fromIntegral (up - low) * (1.0 - exp(-x)))

-- | Print a board to stdout
drawBoard :: Board -> IO ()
drawBoard board = do
  mapM_ (print . col) board
  where
    size = length board
    col :: Int -> [String]
    col r = replicate r " "  <> ["#"] <> replicate (size-r) " "
    
(!) :: (IsSequence seq) => seq -> Index seq -> Element seq
(!) = indexEx