summaryrefslogtreecommitdiff
path: root/16-1.hs
blob: 884f5eccfa116d6a2dec6f6173889ebe4b7fcce9 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
import Control.Monad (filterM, foldM)
import Control.Monad.ST
import Control.Monad.Trans.Class
import Control.Monad.Trans.Maybe
import Data.Array as A
import Data.Array.ST
import Data.Bifunctor (second)
import Data.List as L (find, uncons)
import Data.Maybe (fromJust)
import Data.Tuple as T (swap)
import Data.Vector.Mutable as V hiding (foldM, mapM_)
import Prelude
import Prelude as P (length)

data BucketQueue s a = BQ Int (V.MVector s [a])

instance Show (BucketQueue s a) where show (BQ top _) = show top

insert :: BucketQueue s a -> (Int,a) -> ST s (BucketQueue s a)
insert (BQ top buckets) (prio,x) =
  V.modify buckets (x:) prio >> return (BQ (min top prio) buckets)

view :: BucketQueue s a -> MaybeT (ST s) (a, BucketQueue s a)
view (BQ top buckets) = do
  (x, xs) <- V.read buckets top >>= hoistMaybe . L.uncons
  lift $ V.write buckets top xs
  top' <- lift $ findTop top buckets
  return (x, BQ top' buckets)
{-# SCC view #-}

findTop :: Int -> V.MVector s [a] -> ST s Int
findTop n v = do
  top <- V.read v n
  case top of
    [] -> if n < V.length v - 1 then findTop (n+1) v else return n
    _  -> return n

newBQ :: Int -> ST s (BucketQueue s a)
newBQ n = BQ n <$> V.replicate n []

next :: Direction -> (Int,Int) -> (Int,Int)
next North (x,y) = (x,y-1)
next South (x,y) = (x,y+1)
next West  (x,y) = (x-1,y)
next East  (x,y) = (x+1,y)

left :: Direction -> Direction
left North = West
left West  = South
left South = East
left East  = North

right :: Direction -> Direction
right North = East
right East  = South
right South = West
right West  = North

type Node = ((Int,Int), Direction)

type Frontier s = BucketQueue s Node

data Direction = North | South | East | West deriving (Eq,Ix,Ord,Show)

data Tile = Wall | Tile { untile :: [Direction] }

type Map = A.Array (Int,Int) Tile

type Costs s = STArray s Node (Maybe Int)

neighbours :: (Int,Int) -> [(Direction, (Int,Int))]
neighbours (x,y) = [ (North, (x,y-1))
                   , (South, (x,y+1))
                   , (West,  (x-1,y))
                   , (East,  (x+1,y))
                   ]

parseMap :: [String] -> Map
parseMap s =
  let w = P.length (head s) - 1
      h = P.length s - 1
      f p = case g p of
        '#' -> Wall
        _   -> Tile . map fst . filter ((/= '#') . g . snd) $ neighbours p
      g (x,y) = s !! y !! x
  in array ((0,0),(w,h)) [ ((x,y), f (x,y))
                         | x <- [0..w]
                         , y <- [0..h]
                         ]

findStart :: [String] -> (Int,Int)
findStart ls =
  let Just p = L.find ((== 'S') . snd) [ ((x,y), ls !! y !! x)
                                       | x <- [0..P.length (head ls) - 1]
                                       , y <- [0..P.length ls - 1]
                                       ]
  in fst p

findGoal :: [String] -> (Int,Int)
findGoal ls =
  let Just p = L.find ((== 'E') . snd) [ ((x,y), ls !! y !! x)
                                       | x <- [0..P.length (head ls) - 1]
                                       , y <- [0..P.length ls - 1]
                                       ]
  in fst p

distance :: (Int,Int) -> (Int,Int) -> Int
distance (x0,y0) (x1,y1) =
  let a = fromIntegral $ x0 - x1
      b = fromIntegral $ y0 - y1
  in floor $ sqrt (a^2 + b^2)

expand :: Map -> (Node, Maybe Int) -> [(Node,Maybe Int)]
expand m ((pos, dir), Just cost) =
  [((next dir pos, dir), Just $ cost + 1) | dir `elem` untile (m A.! pos)]
  ++ ((\x -> ((pos,x), Just $ cost + 1000)) <$> [left dir, right dir])

pathThrough :: Map -> (Node,Maybe Int) -> Node -> (Node,Maybe Int)
pathThrough m through dest =
  head . filter ((== dest) . fst) $ expand m through

betterPath :: Map
           -> (Node,Maybe Int)
           -> (Node,Maybe Int)
           -> (Node,Maybe Int)
betterPath m through dest@(node,oldCost) =
  let n@(_,newCost) = pathThrough m through node
  in if newCost < oldCost
     then n
     else dest

cost :: Costs s -> Node -> ST s (Node, Maybe Int)
cost costs node = (node,) <$> readArray costs node

data PathingResult s = Success Int | Fail | Working (Frontier s, Costs s)

heuristic :: (Int,Int) -> (Node,Int) -> (Node,Int)
heuristic goal (n@(p,_),f) = (n, f + distance goal p)

(///) :: (MArray a e m, Ix i, Show i, Show e) => a i e -> [(i,e)] -> m ()
a /// asss = mapM_ (uncurry (writeArray a)) asss

unvisited :: Costs s -> Node -> ST s Bool
unvisited costs node = (== Nothing) <$> readArray costs node

pathOnce :: Map
         -> (Int,Int)
         -> Frontier s
         -> Costs s
         -> ST s (PathingResult s)
pathOnce m goal frontier costs =
  runMaybeT (view frontier) >>= maybe (return Fail) continue
  where
    continue (current@(p,_), rest)
      | p == goal = Success . fromJust <$> readArray costs current
      | otherwise = do
          curCost <- cost costs current
          let nexts = expand m curCost
          frontier' <- filterM (unvisited costs . fst) nexts
            >>= foldM insert rest
            . map (T.swap . heuristic goal . second fromJust)
          costs /// map (betterPath m curCost) nexts
          return . Working $ (frontier', costs)

findPath :: Map -> (Int,Int) -> (Int,Int) -> ST s Int
findPath m start goal = do
    frontier <- {-# SCC "frontier" #-}
      let (_,(w,h)) = bounds m
      in newBQ (100*w*h) >>= flip insert (0,(start,East))

    costs <- let (lower,upper) = bounds m
             in newArray ((lower,North),(upper,West)) Nothing

    writeArray costs (start,East) (Just 0)

    loop frontier costs
      where
        loop :: Frontier s -> Costs s -> ST s Int
        loop frontier costs = do
          once <- pathOnce m goal frontier costs
          case once of
            Success x -> return x
            Fail -> error "no path"
            Working (frontier', costs') -> loop frontier' costs'

main :: IO ()
main = do
  input <- lines <$> getContents
  let m = parseMap input
  let start = findStart input
  let goal = findGoal input
  print (runST $ findPath m start goal :: Int)