Leo Zhang
Leo Zhang

Reputation: 3230

Generic Depth First Search algorithm with Set

I'm trying to implement a generic Depth First Search algorithm with the following type signature:

dfs :: (Ord a) => (a -> [a]) -> a -> [a]

I had a reference from this blog post:

dfs2 :: (Ord a) => (a -> Set a) -> a -> [a]
dfs2 succ start = loop [start] (Set.singleton start) where
    loop [] _ = []
    loop (x:xs) visited = x : loop (Set.toList new ++ xs) (Set.union visited new)
        where new = Set.difference (succ x) visited

It works that it does DFS, and only visit the same element once.

f 1 = Set.fromList [2, 3]
f 2 = Set.fromList [3, 4, 5, 6]
f _ = Set.empty

ghci> dfs2 f 1
[1,2,4,5,6,3]

But since dfs2 takes f :: (a -> Set a) which returns a Set a rather than a [a], I can't specify the order of the elements to visit.

For example, if f2 is defined like this:

f2 1 = Set.fromList [3, 2]
f2 2 = Set.fromList [6, 5, 4, 3]
f2 _ = Set.empty

It would return the same result

ghci> dfs2 f2 1
[1,2,4,5,6,3]

But what I want is [1,3,2,6,5,4]

I can't figure out how to make changes to dfs2 to implement with the type signature

dfs :: (Ord a) => (a -> [a]) -> a -> [a]

such that, dfs would visit each element only once in the order that is specified by [a]. And

dfs f2 1 == [1,3,2,6,5,4]

Does anyone have some idea?

Upvotes: 1

Views: 185

Answers (1)

Cirdec
Cirdec

Reputation: 24166

You can write dfs2 in terms of a set of visited nodes coming from the left and the remainder of the results coming from the right.

My head is stuck in folds this week, so first I'm going to define the idea of a fold that operates on both strict left-associated values coming from the left and lazy right-associated values coming from the right; it'd probably be more illustrative to you to write out an explicit version of dfs2.

{-# LANGUAGE BangPatterns #-}

foldboth :: ((l, r) -> a -> (l, r)) -> (l, r) -> [a] -> (l, r)
foldboth f = go
  where
    go (l, r) [] = (l, r)
    go (l, r) (x:xs) = (l'', r'')
      where
        (l', r'') = f (l, r') x
        (l'', r') = go (l', r) xs

foldboth' :: ((l, r) -> a -> (l, r)) -> (l, r) -> [a] -> (l, r)
foldboth' f = foldboth f'
  where
    f' (!l, r) = f (l, r)

Depth first search can be defined by folding up the successors, skipping the ones that have already been visited and adding the unvisited ones both the the set traveling left-to-right and to the results traveling right to left.

dfs2 :: (Ord a) => (a -> [a]) -> a -> [a]
dfs2 succ start = snd $ go (Set.empty, []) [start] where
  go = foldboth' step
  step (visited, remaining) x = 
    if Set.member x visited
    then (visited, remaining)
    else (visited', x:remaining')
      where
        (visited', remaining') = go (Set.insert x visited, remaining) (succ x) 

This gives the desired result

> dfs2 f2 1
[1,3,2,6,5,4]

Upvotes: 3

Related Questions