Dfs.hs

module Dfs where

import Control.Monad.State
import Data.IntMap (IntMap)
import Data.IntSet (IntSet)
import Data.Maybe (fromMaybe)
import qualified Data.IntMap as IntMap
import qualified Data.IntSet as IntSet

type Vertex = Int
type Graph = IntMap [Vertex]

dfs :: Vertex -> Vertex -> Graph -> Bool
dfs from to edges = evalState (dfsInner [from]) IntSet.empty
  where
    dfsInner :: [Vertex] -> State IntSet Bool
    dfsInner [] = return False
    dfsInner (hd : tl)
      | hd == to = return True
      | otherwise =
          do modify (IntSet.insert hd)
             currentVisited <- get
             let nextToVisit = filter (not . flip IntSet.member currentVisited)
                               (fromMaybe [] (IntMap.lookup hd edges))
             dfsInner (nextToVisit ++ tl)

testGraph1 :: Graph
testGraph1 = IntMap.fromList
  [ (1, [2, 3])
  , (2, [1, 3, 4])
  , (3, [1, 2])
  , (4, [2])
  ]

testGraph2 :: Graph
testGraph2 = IntMap.fromList
  [ (1, [2, 3])
  , (2, [1, 3])
  , (3, [1, 2])
  , (4, [])
  ]