Project Euler 96 - Solve Sudoku! (in Haskell)

sudoku.hs

#!/usr/bin/env runhaskell

{- Project Euler 96 - Solve Sudoku!
 -
 - Compile it with ghc -o blah --make sudoku
 -
 - And grab sudoku.txt from 
 - http://projecteuler.net/index.php?section=problems&id=96
 -
 - Answer: 24702
 -}

import Debug.Trace

-- There is probably a better way of doing this but 
-- why not try to do it myself...
newlineIterator :: String -> [String]
newlineIterator s = go [] "" s
    where
        go :: [String] -> String -> String -> [String]
        go ls "" "" = ls
        go ls cs "" = ls ++ [cs]
        go ls cs (c:'\r':'\n':end) = go (ls ++ [cs++[c]]) "" end 
        go ls cs (c:end) = go ls (cs++[c]) end

newtype Grid = WrapGrid [[Int]] deriving (Show)
makeGrid :: [String] -> Grid
makeGrid rs = WrapGrid [[read [i] | i <- is] | is <- rs]

-- get row i from a grid
row :: Grid -> Int -> [Int]
row (WrapGrid rs) i | (i < 0) || (i > 8) = error "invalid row"
                    | otherwise = rs !! i

-- get col i from a grid
col :: Grid -> Int -> [Int]
col (WrapGrid rs) i | (i < 0) || (i > 8) = error "invalid col"
                    | otherwise = [row !! i | row <- rs]

flatten :: [[a]] -> [a]
flatten xs = go [] xs
    where
        go ys [] = ys
        go ys (x:xs) = go (ys ++ x) xs

-- get block i from a grid (where block 0 is the top left square, 
-- block 1 is the top middle, and so on)
block :: Grid -> Int -> [Int]
block (WrapGrid rs) i | (i < 0) || (i > 8) = error "invalid block"
                      | otherwise = flatten rs' 
                      where 
                        f k = (take 3) . (drop (k * 3))
                        rs' = [(f mj) r | r <- (f mi) rs]
                        mi = div i 3
                        mj = mod i 3

-- are all 9 numbers present?
segCorrect :: [Int] -> Bool
segCorrect row = go row 1
    where
        go :: [Int] -> Int -> Bool
        go _ 10 = True
        go row i | elem i row = go row (i+1)
                 | otherwise = False

firstIndex :: (a -> Bool) -> [a] -> Maybe Int
firstIndex p xs = go p xs 0
    where 
        go _ [] _ = Nothing
        go p (x:xs) i | p x = Just i
                      | otherwise = go p xs (i+1)

first :: (a -> Bool) -> [a] -> Maybe a
first p [] = Nothing
first p (x:xs) | p x = Just x
               | otherwise = first p xs

-- pluck the next open row index out of the
-- grid... first row with an empty space
nextOpenRowIndex :: Grid -> Maybe Int
nextOpenRowIndex grid = 
    let
        has0 :: [Int] -> Bool
        has0 = elem 0
        rows :: [[Int]]
        rows = (map (row grid) [0..8])
    in firstIndex (has0) rows

-- pluck the first open space coords out of the grid
nextOpenRowColIndices :: Grid -> Maybe (Int, Int)
nextOpenRowColIndices grid = do 
    rowIndex <- nextOpenRowIndex grid
    let targetRow = (row grid rowIndex)
    colIndex <- firstIndex (== 0) targetRow
    return (rowIndex, colIndex)
      
-- should use Data.Set
intersection :: (Eq a) => [a] -> [a] -> [a]
intersecion _ [] = []
intersecion [] _ = []
intersection (x:xs) ys | elem x ys = x : (intersection xs ys)
                       | otherwise = intersection xs ys
intersection _ _ = []

-- is our puzzle completely filled in?
filledIn :: Grid -> Bool
filledIn grid = (Nothing == nextOpenRowIndex grid)

-- is it filled in correctly?
gridCorrect :: Grid -> Bool
gridCorrect grid = (filledIn grid) && 
                   (all id (map segCorrect (map (row grid) [0..8]))) &&
                   (all id (map segCorrect (map (col grid) [0..8]))) &&
                   (all id (map segCorrect (map (block grid) [0..8])))

-- utility for parsing a file into grids
gridIterator :: [String] -> [Grid]
gridIterator ls = go [] [] ls
    where
        go :: [Grid] -> [String] -> [String] -> [Grid]
        go gs [] [] = gs
        go gs ls [] | length ls == 10 = gs ++ [makeGrid $ tail ls]
                    | otherwise       = error "Invalid grid len"
        go gs ls (r:rs) | length ls == 10 = go (gs ++ [makeGrid $ tail ls]) [] (r:rs)
                        | otherwise       = go gs (ls ++ [r]) rs

-- what can we fill into an empty space?
getMissing :: [Int] -> [Int]
getMissing row = go row [] 1
    where
        go row acc 10 = acc
        go row acc i | elem i row = go row acc (i+1)
                     | otherwise = go row (acc ++ [i]) (i+1)

debug = flip trace

-- destructuring and restructuring lists... not the best
fillIn :: [Int] -> Int -> [Int]
fillIn row m = (go m [] row) --`debug` ("fillIn " ++ (show row) ++ " " ++ (show m)) 
    where
        go m acc [] = acc
        go m acc (x:xs) | x == 0 = acc ++ [m] ++ xs
                        | otherwise = go m (acc ++ [x]) xs 

replaceRow :: Grid -> [Int] -> Int -> Grid
replaceRow (WrapGrid gridRows) row index = go row index 0 [] gridRows
    where
        go row index curIndex acc [] = WrapGrid acc
        go row index curIndex acc (r:rs) | index == curIndex = go row index (curIndex+1) (acc ++ [row]) rs
                                         | otherwise = go row index (curIndex+1) (acc ++ [r]) rs

getBlockIndex :: Int -> Int -> Int
getBlockIndex r c = 3*(div r 3) + (div c 3)

-- grid solution candidates
enumerate :: Grid -> [Grid]
enumerate grid = case maybeIndices of
                Nothing -> [grid]
                Just (rowIndex, colIndex) -> do
                    let {origRow = row grid rowIndex;
                         origCol = col grid colIndex;
                         blockIndex = getBlockIndex rowIndex colIndex;
                         origBlock = block grid blockIndex;
                        rowMissing = getMissing origRow;
                        colMissing = getMissing origCol;
                        blockMissing = getMissing origBlock;
                        missing = intersection (intersection rowMissing colMissing) blockMissing}
                    m <- missing
                    let {newRow = fillIn origRow m;
                        grid' = replaceRow grid newRow rowIndex}
                    (enumerate grid')-- `debug` (show grid')
    where 
        maybeIndices = nextOpenRowColIndices grid

-- pluck the first correct solution out of the candidates
solve :: Grid -> Grid
solve grid = case solution of
                Nothing -> error "unsolvable problem"
                Just wrappedSoln -> wrappedSoln
    where solution = first gridCorrect (enumerate grid)

printI :: Grid -> Int -> IO()
printI g i = do    
    print $ "ROW "++(show i)
    print $ row g i
    print $ "COL "++(show i)
    print $ col g i
    print $ "BLK "++(show i)
    print $ block g i

dot :: [Int] -> [Int] -> Int
dot (x:xs) (y:ys) = (x*y) + (dot xs ys) 
dot [] [] = 0
dot _ _ = error "Invalid len"

-- this is the magic number the problem wants us to calculate
derivedValue :: Grid -> Int
derivedValue (WrapGrid grid) = dot [100,10,1] (take 3 $ head grid)

main = do
    grids <- (readFile "sudoku.txt") >>= return.newlineIterator >>= return . gridIterator 
    --let g = grids !! 0
    --print g
    --mapM_ (printI g) [0..8]
    --let g' = solve g
    --print g'
    --print $ gridCorrect g'
    --print $ derivedValue g'
    print "thinking..."
    print $ sum (map derivedValue (map solve grids))