src/task8/task8.hs

   1 import Data.Maybe (fromMaybe)
   2 import qualified Data.PQueue.Prio.Min as P
   3 import qualified Data.HashSet as S
   4 import qualified Data.Sequence as Q
   5 import qualified Data.HashMap.Strict as M
   6 import Data.Sequence ((<|))
   7 import Data.Foldable (foldl')
   8 import Data.HashMap.Strict ((!))
   9
  10
  11 data Grass = Short | Long deriving (Show, Eq)
  12 type Location = (Int, Int)
  13 data Neighbour = Neighbour {stepLocation :: Location, stepCost :: Int} deriving (Show, Eq)
  14 type Lawn = M.HashMap Location Grass
  15 type Closed = S.HashSet Location
  16 data Agendum = Agendum {current :: Location, trail :: Q.Seq Location, cost :: Int} deriving (Show, Eq)
  17 type Agenda = P.MinPQueue Int Agendum
  18
  19
  20 main :: IO ()
  21 main = do
  22     lawnT <- readFile "data/08-maze.txt"
  23     let lawn = parseLawn lawnT
  24     let goal = goalLocation lawn
  25     print $ part1 lawn goal
  26     print $ part2 lawn goal
  27
  28 part1 :: Lawn -> Location -> Int
  29 part1 lawn goal = cost $ fromMaybe (snd $ P.findMin $ initAgenda goal) $ aStar 999 goal lawn (initAgenda goal) S.empty
  30
  31 part2 :: Lawn -> Location -> Int
  32 part2 lawn goal = cost $ fromMaybe (snd $ P.findMin $ initAgenda goal) $ aStar 3 goal lawn (initAgenda goal) S.empty
  33
  34
  35 parseLawn :: String -> Lawn
  36 parseLawn lawnT = M.fromList lawnCells
  37     where lawnLines = zip [1..] $ lines lawnT
  38           lawnCells = concatMap (\(r, row) -> zipWith (cellify r) [1..] row) lawnLines
  39           cellify r c cell
  40             | cell == '#' = ((r, c), Long)
  41             | otherwise   = ((r, c), Short)
  42
  43 goalLocation :: Lawn -> Location
  44 goalLocation lawn = ( maximum $ map fst $ M.keys lawn
  45                     , maximum $ map snd $ M.keys lawn
  46                     )
  47
  48 initAgenda :: Location -> Agenda
  49 initAgenda goal = P.singleton (estimateCost (1, 1) goal) Agendum {current = (1, 1), trail = Q.empty, cost = 0}
  50
  51
  52 aStar :: Int -> Location -> Lawn -> Agenda -> Closed -> Maybe Agendum
  53 aStar longCost goal lawn agenda closed
  54     -- | trace ("Peeping " ++ (show $ fst $ P.findMin agenda) ++ ": " ++ (show reached) ++ " <- " ++ (show $ toList $ Q.take 1 $ trail $ currentAgendum) ++ " :: " ++ (show newAgenda)) False = undefined
  55     | P.null agenda = Nothing
  56     | otherwise =
  57         if reached == goal then Just currentAgendum
  58         else if reached `S.member` closed
  59             then aStar longCost goal lawn (P.deleteMin agenda) closed
  60             else aStar longCost goal lawn newAgenda (S.insert reached closed)
  61         where
  62             (_, currentAgendum) = P.findMin agenda
  63             reached = current currentAgendum
  64             newAgenda = foldl' (\q a -> P.insert (estimatedCost a) a q) (P.deleteMin agenda) $ candidates longCost lawn currentAgendum closed
  65             estimatedCost agendum = estimateCost (current agendum) goal + cost agendum
  66
  67
  68
  69 candidates :: Int -> Lawn -> Agendum -> Closed -> Q.Seq Agendum
  70 candidates longCost lawn agendum closed = newCandidates
  71     where
  72         candidate = current agendum
  73         previous = trail agendum
  74         succs = successors longCost lawn candidate
  75         nonloops = Q.filter (\s -> not $ (stepLocation s) `S.member` closed) succs
  76         newCandidates = fmap (\n -> makeAgendum n) nonloops
  77         makeAgendum new = Agendum {current = stepLocation new,
  78                                     trail = candidate <| previous,
  79                                     cost = cost agendum + stepCost new}
  80
  81
  82 successors :: Int -> Lawn -> Location -> (Q.Seq Neighbour)
  83 successors longCost lawn (row, column) = Q.fromList $ map neighbourify neighbours
  84     where neighbours = filter (\l -> l `M.member` lawn)
  85                                 [(r, c) | r <- [(row - 1)..(row + 1)],
  86                                           c <- [(column - 1)..(column + 1)],
  87                                           r >= minR,
  88                                           r <= maxR,
  89                                           c >= minC,
  90                                           c <= maxC,
  91                                           ((r == row && c /= column) || (r /= row && c == column)) ]
  92           neighbourify neighbour = Neighbour {stepLocation = neighbour, stepCost = scCalc neighbour}
  93           minR = minimum $ map fst $ M.keys lawn
  94           maxR = maximum $ map fst $ M.keys lawn
  95           minC = minimum $ map snd $ M.keys lawn
  96           maxC = maximum $ map snd $ M.keys lawn
  97           scCalc location = if lawn!location == Long
  98                             then longCost
  99                             else 1
 100
 101
 102 estimateCost :: Location -> Location -> Int
 103 estimateCost (r, c) (gr, gc) = abs (r - gr) + abs(c - gc)
 104