Haskell Quiz/Word Search/Solution Sjanssen

A trie structure

A prefix oriented structure is ideal for the word search algorithm.

module Trie where

import qualified Data.Map as Map
data Trie = Trie Bool (Map.Map Char Trie) deriving (Show)

empty = Trie False Map.empty

insert []     (Trie b m) = Trie True m
insert (x:xs) (Trie b m) = Trie b $ Map.alter (maybe (Just $ fromString xs) (Just . insert xs)) x m

fromString = foldr (\x xs -> Trie False (Map.singleton x xs)) (Trie True Map.empty)

fromList = foldr insert empty

Initial solution

import Data.Array
import Data.Char
import qualified Data.Map as Map

import Trie

type Grid = Array (Int, Int) Char

-- | Given a grid and a trie, produce a list of matched words as well as the
-- location of the match.  The algorithm differs slightly from the example,
-- if a word appears several times a match is produced for each location.
match :: Grid -> Trie -> [(String, [(Int, Int)])]
match g t = concatMap (flip go t) paths
 where
    bs     = bounds g
    paths  = concatMap dirs (range bs)
    dirs x = [takeWhile (inRange bs) $ iterate (next (i, j)) x
                    | i <- [-1 .. 1], j <- [-1 .. 1], i /= 0 || j /= 0]
    go xs (Trie b m) = (if b then [([], [])] else [])
                        ++ do
                            (x:xs) <- return xs
                            let c = g ! x
                            t <- Map.lookup c m
                            (s, is) <- go xs t
                            return (c:s, x:is)

next (x, y) (i, j) = (i + x, j + y)

readBoard = do
    l <- getLine
    if null l
        then return []
        else fmap (l:) readBoard

csv [] = []
csv xs = l : csv (dropWhile (\x -> isSpace x || x == ',') r)
 where (l, r) = span (/= ',') xs

output g ms = sequence_ [putStrLn [g' ! (i, j) | j <- [1..c]] | i <- [1..r]]
 where
    g' = listArray ((1, 1), (r, c)) (repeat '+') // (map (\x -> (x, g ! x)) $ concatMap snd ms)
    ((1, 1), (r, c)) = bounds g

main = do
    b <- readBoard
    let r = length b
        c = length (head b)
        g = listArray ((1, 1), (r, c)) . map toUpper . concat $ b
    l <- getLine
    let ws = map (map toUpper) $ csv l
        t  = fromList ws
        ms = match g t
    output g ms

Extra credit

For extra credit, I extended the program to support wildcards (*) in the board. This only required changing two lines and adding another.

import Data.Array
import Data.Char
import qualified Data.Map as Map

import Trie

type Grid = Array (Int, Int) Char

-- | Given a grid and a trie, produce a list of matched words as well as the
-- location of the match.  The algorithm differs slightly from the example,
-- if a word appears several times a match is produced for each location.
match :: Grid -> Trie -> [(String, [(Int, Int)])]
match g t = concatMap (flip go t) paths
 where
    bs     = bounds g
    paths  = concatMap dirs (range bs)
    dirs x = [takeWhile (inRange bs) $ iterate (next (i, j)) x
                    | i <- [-1 .. 1], j <- [-1 .. 1], i /= 0 || j /= 0]
    go xs (Trie b m) = (if b then [([], [])] else [])
                        ++ do
                            (x:xs) <- return xs
                            (c, t) <- case g ! x of
                                        '*' -> Map.assocs m
                                        c   -> fmap ((,) c) (Map.lookup c m)
                            (s, is) <- go xs t
                            return (c:s, x:is)

next (x, y) (i, j) = (i + x, j + y)

readBoard = do
    l <- getLine
    if null l
        then return []
        else fmap (l:) readBoard

csv [] = []
csv xs = l : csv (dropWhile (\x -> isSpace x || x == ',') r)
 where (l, r) = span (/= ',') xs

output g ms = sequence_ [putStrLn [g' ! (i, j) | j <- [1..c]] | i <- [1..r]]
 where
    g' = listArray ((1, 1), (r, c)) (repeat '+') // (map (\x -> (x, g ! x)) $ concatMap snd ms)
    ((1, 1), (r, c)) = bounds g

main = do
    b <- readBoard
    let r = length b
        c = length (head b)
        g = listArray ((1, 1), (r, c)) . map toUpper . concat $ b
    l <- getLine
    let ws = map (map toUpper) $ csv l
        t  = fromList ws
        ms = match g t
    output g ms