Haskell Quiz/Astar/Solution Dolio

A* requires keeping a priority queue of places to visit. This can be done with a simple sorted list, but I decided to make a PriorityQueue data type for use in the algorithm instead. The implementation uses lazy pairing heaps from Chris Okasaki's Purely Functional Data Structures.

module PriorityQueue (
        PriorityQueue,
        empty,
        singleton,
        fromList,
        null,
        deleteFindMin,
        deleteMin,
        findMin,
        insert,
        union
       ) where

import Prelude hiding (null)
import Data.List (foldl')

data Ord k => PriorityQueue k a = Nil | Branch k a (PriorityQueue k a) (PriorityQueue k a)

empty :: Ord k => PriorityQueue k a
empty = Nil

singleton :: Ord k => k -> a -> PriorityQueue k a
singleton k a = Branch k a Nil Nil

fromList :: Ord k => [(k,a)] -> PriorityQueue k a
fromList = foldl' (\q (k,a) -> singleton k a `union` q) empty

null :: Ord k => PriorityQueue k a -> Bool
null Nil = True
null _   = False

deleteFindMin :: Ord k => PriorityQueue k a -> ((k,a), PriorityQueue k a)
deleteFindMin Nil = error "Empty heap."
deleteFindMin (Branch k a l r) = ((k,a), union l r)

deleteMin :: Ord k => PriorityQueue k a -> PriorityQueue k a
deleteMin h = snd (deleteFindMin h)

findMin :: Ord k => PriorityQueue k a -> (k, a)
findMin h = fst (deleteFindMin h)

insert :: Ord k => k -> a -> PriorityQueue k a -> PriorityQueue k a
insert k a h = union (singleton k a) h

union :: Ord k => PriorityQueue k a -> PriorityQueue k a -> PriorityQueue k a
union l Nil = l
union Nil r = r
union l@(Branch kl _ _ _) r@(Branch kr _ _ _)
    | kl <= kr  = link l r
    | otherwise = link r l

link (Branch k a Nil m) r = Branch k a r m
link (Branch k a ll lr) r = Branch k a Nil (union (union r ll) lr)

Not all the functions from data structures in the standard library (Data.Map, Data.Set, etc.) are provided; I only wrote those that are needed for the algorithm. However, this could be extended easily.

The rest is just a general A* function, which takes a starting place, and functions for successors, testing for completion, cost of a place, and heuristic estimation from a place to the end, returning the path taken (a list from end to start). The rest of the code deals with the specifics of the ASCII map:

{-# OPTIONS_GHC -fglasgow-exts #-}

module Main where
import Control.Monad (guard, liftM2)
import Control.Monad.Instances
import Data.List (elemIndex)
import qualified Data.Set as S
import qualified Data.Map as M
import qualified PriorityQueue as Q

type Point = (Int, Int)
type Map = [[Char]]

find :: Char -> Map -> Point
find c m = find' 0 m
 where find' _ [] = error "Can't find tile."
       find' y (h:t)
        | Just x <- elemIndex c h = (y, x)
        | otherwise = find' (y+1) t

heuristic :: Point -> Point -> Int
heuristic (x, y) (u, v) = abs (x - u) `max` abs (y - v)

successor :: Map -> Point -> [Point]
successor m (x,y) = do u <- [x + 1, x, x - 1]
                       v <- [y + 1, y, y - 1]
                       guard (0 <= u && u < length m)
                       guard (0 <= v && v < length (head m))
                       guard (u /= x || y /= v)
                       guard (m !! u !! v /= '~')
                       return (u, v)

astar start succ end cost heur 
    = astar' (S.singleton start) (Q.singleton (heur start) [start])
 where
 astar' seen q
    | Q.null q  = error "No Solution."
    | end n     = next
    | otherwise = astar' seen' q'
  where
  ((c,next), dq) = Q.deleteFindMin q
  n     = head next
  succs = filter (`S.notMember` seen) $ succ n
  costs = map ((+ c) . (subtract $ heur n) . liftM2 (+) cost heur) succs
  q'    = dq `Q.union` Q.fromList (zip costs (map (:next) succs))
  seen' = seen `S.union` S.fromList succs

path :: [[Char]] -> [Point] -> [[Char]]
path m l = iterY m l 0
 where iterY [] _ _ = []
       iterY (h:t) l n = iterX h l n 0 : iterY t l (n+1)
       iterX [] _ _ _ = []
       iterX (h:t) l n m = (if (n,m) `elem` l then '#' else h) : iterX t l n (m+1)

doit s = unlines . path m $ astar start succ (== end) cost h
 where m     = lines s
       start = find '@' m
       end   = find 'X' m
       succ  = successor m
       h     = heuristic end
       cost (x, y) = costsM M.! (m !! x !! y)
       costsM = M.fromList [('@',1),('x',1),('X',1),('.',1),('*',2),('^',3)]

main = interact doit