Tree Algorithms in Haskell

TreeAlgorithms.hs

{-# LANGUAGE NoImplicitPrelude #-}

module TreeAlgorithms where

import ClassyPrelude

main :: IO ()
main = do
  let tree = binTree 5    -- create binary tree with depth 5
  print (dfs tree 1)      -- list nodes with depth first search
  print (bfs tree 1)      -- list nodes with breadth first search
  print (dls tree 4 1 31) -- does 31 exist max 4 levels deep?
  print (idfds tree 1 31) -- find 31 and show depth if found

type Node = Int
type Tree = Map Node [Node]

-- | Binary Tree
binTree :: Int -> Tree
binTree depth = mapFromList [ (i, [i*2,i*2+1]) | i <- [1..2^(depth-1) - 1] ]

-- | depth first search
dfs :: Tree -> Node -> [Node]
dfs tree v = v : concatMap (dfs tree) (findWithDefault [] v tree)

-- | breadth first search
bfs :: Tree -> Node -> [Node]
bfs tree v = _bfs [v]
  where
    _bfs [] = []
    _bfs (v:vs) = v : _bfs (vs ++ (findWithDefault [] v tree))

-- | depth limited search (depth first)
dls :: Tree -> Int -> Node -> Node -> Bool
dls tree depth v w
  | w == v    = True
  | depth < 2 = False
  | otherwise = any (\c -> dls tree (depth-1) c w) children
  where
    children = findWithDefault [] v tree

-- | iterative (depth first) deepening search
idfds :: Tree -> Node -> Node -> Maybe Int -- depth of wanted node
idfds tree v w = find (\i -> dls tree i v w) maxDepth
  where
    maxDepth = [1..1000] :: [Int]