76 lines
3.1 KiB
Haskell
76 lines
3.1 KiB
Haskell
-- https://gist.github.com/abhin4v/8172534
|
|
|
|
module AStar where
|
|
|
|
import qualified Data.PQueue.Prio.Min as PQ
|
|
import qualified Data.HashSet as Set
|
|
import qualified Data.HashMap.Strict as Map
|
|
import Data.Hashable (Hashable)
|
|
import Data.List (foldl')
|
|
import Data.Maybe (fromJust)
|
|
|
|
{- cabal:
|
|
build-depends: base, pqueue, unordered-containers, hashable
|
|
-}
|
|
|
|
-- A* search: Finds the shortest path from a start node to a goal node using a heuristic function.
|
|
astarSearch :: (Eq a, Hashable a) =>
|
|
a -- startNode: the node to start the search from
|
|
-> (a -> Bool) -- isGoalNode: a function to test if a node is the goal node
|
|
-> (a -> [(a, Int)]) -- nextNodeFn: a function which calculates the next nodes for a current node
|
|
-- along with the costs of moving from the current node to the next nodes
|
|
-> (a -> Int) -- heuristic: a function which calculates the (approximate) cost of moving
|
|
-- from a node to the nearest goal node
|
|
-> Maybe (Int, [a]) -- result: Nothing is no path is found else
|
|
-- Just (path cost, path as a list of nodes)
|
|
astarSearch startNode isGoalNode nextNodeFn heuristic =
|
|
astar (PQ.singleton (heuristic startNode) (startNode, 0))
|
|
Set.empty (Map.singleton startNode 0) Map.empty
|
|
where
|
|
-- pq: open set, seen: closed set, tracks: tracks of states
|
|
astar pq seen gscore tracks
|
|
-- If open set is empty then search has failed. Return Nothing
|
|
| PQ.null pq = Nothing
|
|
-- If goal node reached then construct the path from the tracks and node
|
|
| isGoalNode node = Just (gcost, findPath tracks node)
|
|
-- If node has already been seen then discard it and continue
|
|
| Set.member node seen = astar pq' seen gscore tracks
|
|
-- Else expand the node and continue
|
|
| otherwise = astar pq'' seen' gscore' tracks'
|
|
where
|
|
-- Find the node with min f-cost
|
|
(node, gcost) = snd . PQ.findMin $ pq
|
|
|
|
-- Delete the node from open set
|
|
pq' = PQ.deleteMin pq
|
|
|
|
-- Add the node to the closed set
|
|
seen' = Set.insert node seen
|
|
|
|
-- Find the successors (with their g and h costs) of the node
|
|
-- which have not been seen yet
|
|
successors =
|
|
filter (\(s, g, _) ->
|
|
not (Set.member s seen') &&
|
|
(not (s `Map.member` gscore)
|
|
|| g < (fromJust . Map.lookup s $ gscore)))
|
|
$ successorsAndCosts node gcost
|
|
|
|
-- Insert the successors in the open set
|
|
pq'' = foldl' (\q (s, g, h) -> PQ.insert (g + h) (s, g) q) pq' successors
|
|
|
|
gscore' = foldl' (\m (s, g, _) -> Map.insert s g m) gscore successors
|
|
|
|
-- Insert the tracks of the successors
|
|
tracks' = foldl' (\m (s, _, _) -> Map.insert s node m) tracks successors
|
|
|
|
-- Finds the successors of a given node and their costs
|
|
successorsAndCosts node gcost =
|
|
map (\(s, g) -> (s, gcost + g, heuristic s)) . nextNodeFn $ node
|
|
|
|
|
|
-- Constructs the path from the tracks and last node
|
|
findPath tracks node =
|
|
if Map.member node tracks
|
|
then findPath tracks (fromJust . Map.lookup node $ tracks) ++ [node]
|
|
else [node] |