Attack-Rooks challenge [closed]

Question

Attack-Rooks is a challenge proposed by ACM-ICPC as a variant of the N-Queens problem, using Rooks instead of Queens. Of course, to disallow the trivial "fill the diagonal" solution, pawns are used to block the rooks from capturing each other.

Given a NxN (N > 0) chess board containing some pawns, output the highest possible number of rooks that can be put in this board without any rook being able to capture each other.

Here is my solution:

Rooks.hs

{-|
Module: Rooks
Description: solution for the Attack Rooks problem

Provides a solution for the Attack Rooks problem (N-Pawn/Rook problem).
-}
module Rooks
    ( Piece(..)
    , Board
    , highestRooks
    ) where

import Data.Matrix
import qualified Data.Vector as Vec

-- | Define a Piece on the board
data Piece = Pawn -- ^ A blocking Pawn
           | Rook -- ^ A placed Rook
           | Empty -- ^ An empty slot
            deriving (Show, Eq)

-- | Defines a Board (synonym to Matrix Piece)
type Board = Matrix Piece

-- | Solves the problem for a given board
highestRooks :: Board -> Int
highestRooks b = rookCount b + highestRooks' b


highestRooks' :: Board -> Int
highestRooks' = (+1) . vecmax . Vec.map highestRooks' . derivedStates
    where vecmax v | Vec.null v = (-1)
                   | otherwise  = Vec.maximum v

rookCount :: Board -> Int
rookCount = Vec.length . Vec.filter (== Rook) . getMatrixAsVector

derivedStates :: Board -> Vec.Vector Board
derivedStates b = Vec.map (placeRook b) . Vec.filter (canPutRook b) $ positions b

positions :: Board -> Vec.Vector (Int, Int)
positions b = Vec.map (\(x,y) -> (x+1, y+1)) $ Vec.generate (n*n)
                (\i -> (i `mod` n, i `div` n))
                where n = nrows b

placeRook :: Board -> (Int, Int) -> Board
placeRook b p = setElem Rook p b

canPutRook :: Board -> (Int, Int) -> Bool
canPutRook b (row, col) = canPutRookInLine (getCol col b) row
                       && canPutRookInLine (getRow row b) col

canPutRookInLine :: Vec.Vector Piece -> Int -> Bool
canPutRookInLine line pos = (Vec.foldr (step pos)
                                       MayPut
                                       (Vec.zip (Vec.generate 3 (+1))  line)
                            ) `elem` [Ok, Put]

data FoldState = MayPut | CannotPut | Put | Ok | Invalid deriving (Eq)

step :: Int -> (Int, Piece) -> FoldState -> FoldState
step pos (i, Pawn) MayPut = if pos == i then Invalid else MayPut
step pos (i, Rook) MayPut = CannotPut
step pos (i, Empty) MayPut = if pos == i then Put else MayPut
step pos (i, Pawn) CannotPut = MayPut
step pos (i, Empty) CannotPut = if pos == i then Invalid else CannotPut
step pos (i, Pawn) Put = Ok
step pos (i, Rook) Put = Invalid
step pos (i, Empty) Put = Put
step pos (i, p) Ok = Ok
step pos (i, p) Invalid = Invalid
step _ _ _ = Invalid -- impossible, but we're making the function total anyway

Main.hs

module Main
    where

import Rooks
import Data.Matrix
import Data.List.Split (splitOn)
import System.Environment (getArgs)
import Control.Monad (forM_)

main :: IO ()
main = getArgs >>= flip forM_ process

process :: String -> IO ()
process = printResult . calculate

printResult :: Maybe Int -> IO ()
printResult Nothing  = putStrLn "Invalid input"
printResult (Just r) = print r

calculate :: String -> Maybe Int
calculate = fmap highestRooks . validate valid . parse

validate :: (a -> Bool) -> Maybe a -> Maybe a
validate _ Nothing = Nothing
validate f (Just x) = check (f x) x
        where   check True  a = Just a
                check False _ = Nothing

valid :: Board -> Bool
valid b = ncols b == nrows b

parse :: String -> Maybe Board
parse = fmap fromLists . parse'

parse' :: String -> Maybe [[Piece]]
parse' = allJust . map parseRow . splitOn "|"

allJust :: [Maybe [a]] -> Maybe [[a]]
allJust [] = Just []
allJust ((Just []):_) = Just []
allJust (Nothing:_) = Nothing
allJust ((Just x):xs) = fmap (x:) (allJust xs)

parseRow :: String -> Maybe [Piece]
parseRow []       = Just []
parseRow ('.':xs) = fmap (Empty :) (parseRow xs)
parseRow ('X':xs) = fmap (Pawn  :) (parseRow xs)
parseRow _        = Nothing

I assumed the input to be given in the format of "...|.X.|...", representing a 3x3 board with a pawn in the center.