a star search

See On Github



Generic placeholder thumbnail

by abhin4v

in haskell

Source Code

module AStarSearch 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)

-- 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
    -- 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'
        -- 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]
{-# LANGUAGE TemplateHaskell, GeneralizedNewtypeDeriving #-}

module AStarSearchTest where

import AStarSearch
import Data.Hashable (Hashable)
import Data.List (foldl')
import Data.Maybe (fromJust, fromMaybe)
import Data.Sequence ((|>))
import Test.QuickCheck
import Test.QuickCheck.All
import qualified Data.HashMap.Strict as Map
import qualified Data.Sequence as Seq

-- We use A* search to find the shortest path (path with least number of moves) of a knight
-- from a start square to a goal square on a chess board.

newtype Square = Square (Int, Int) deriving (Eq, Hashable, Show)
type Path = [Square]

-- Finds the next squares a knight can move to from a given square
nextKnightPos (Square (x, y)) =
  map Square . filter isValidMove . map (\(dx, dy) -> (x + dx, y + dy)) $ moves
    moves = [(1,2), (1,-2), (-1,2), (-1,-2), (2,1), (2,-1), (-2,1), (-2,-1)]
    isValidMove (x, y) = and [x > 0, x < 9, y > 0, y < 9]

-- Creates the heuristic function given a goal square. The heuristic used is half of the max of
-- the distance between x coordinates and the distance between y coordinates.
mkHeuristic (Square (gx, gy)) (Square (x, y)) = max (abs (x-gx)) (abs (y-gy)) `div` 2

-- Finds the shortest path of the knight, returns empty path if the goal is invalid
knightsShortestPath :: Square -> Square -> Path
knightsShortestPath initSq goalSq =
  snd . fromMaybe (0, [])
  $ astarSearch initSq (== goalSq) (flip zip (repeat 1) . nextKnightPos) (mkHeuristic goalSq)

-- Finds the shortest path using breadth first search. Used for checking if the path returned by
-- A* search is indeed shortest.
bfs :: Square -> Square -> Path
bfs startSq goalSq =
  bfs' goalSq (Map.singleton startSq noSuchSq) (Seq.singleton startSq)
    noSuchSq = Square (-1, -1)

    bfs' goalSq tracks open = let
      (first, rest) = Seq.splitAt 1 open
      currentSq = Seq.index first 0
      in if currentSq == goalSq
         then consPath currentSq
         else let
           nextSqs = filter (not . flip Map.member tracks) . nextKnightPos $ currentSq
           tracks' = foldl' (\t s -> Map.insert s currentSq t) tracks nextSqs
           in bfs' goalSq tracks' (foldl (|>) rest nextSqs)
        consPath square =
          if Map.member square tracks
          then consPath (fromJust . Map.lookup square $ tracks) ++ [square]
          else []

-- Setup to generate arbitrary squares for testing
instance Arbitrary Square where
  arbitrary = do
    x <- choose (1, 8)
    y <- choose (1, 8)
    return $ Square (x, y)

-- Properties to test
prop_path_starts_with_start_square startSq goalSq =
  head (knightsShortestPath startSq goalSq) == startSq

prop_path_ends_with_goal_square startSq goalSq =
  last (knightsShortestPath startSq goalSq) == goalSq

prop_path_consists_of_valid_knights_moves startSq goalSq =
  let path = knightsShortestPath startSq goalSq
  in all isValidKnightsMove $ zip path (tail path)
    isValidKnightsMove (sqFrom, sqTo) = sqTo `elem` nextKnightPos sqFrom

prop_no_path_for_invalid_goal startSq =
  knightsShortestPath startSq (Square (-1, -1)) == []

prop_path_is_shortest startSq goalSq =
  length (knightsShortestPath startSq goalSq) == length (bfs startSq goalSq)

-- Tests all the properties 1000 times each
testAllProps = $forAllProperties $ quickCheckWithResult (stdArgs {maxSuccess = 1000})

-- main function runs the tests. Type `runhaskell AStarSearch_test.hs` on command line to run the tests.
main = testAllProps >> return ()