BFS

BFS.hs

{-# LANGUAGE LambdaCase #-}

import Control.Applicative
import Control.Monad.State
import System.Random

data Tree a = Empty | Bin a (Tree a) (Tree a)
    deriving Show

-- O(n)
numAtDepths :: Tree a -> [Int]
numAtDepths = flip execState [] . go
    where go Empty = return ()
          go (Bin _ l r) = get >>= \case
            []     -> go l >> go r >> modify (1:)
            (x:xs) -> put xs >> go l >> go r >> modify (1 + x :)

-- O(n)
label :: Tree a -> Tree (Int, a)
label t = evalState (go (scanl (+) 0 ds) t) $ repeat 0
    where ds = numAtDepths t
          go _ Empty = return Empty
          go (n:ns) (Bin x l r) = get >>= \(m:ms) ->
            Bin (n + m + 1, x) <$> (put ms *> go ns l)
                               <*> go ns r
                               <*  modify (m + 1 :)

randomTree = Bin () <$> subTree <*> subTree
    where subTree = state random >>= \b -> if b then randomTree else return Empty

testT = Bin 'a' 
    (Bin 'b' 
        (Bin 'd'
            Empty 
            (Bin 'h' 
                (Bin 'j' Empty Empty)
                Empty))
        (Bin 'e' Empty Empty))

    (Bin 'c'
        (Bin 'f'
            Empty
            (Bin 'i'
                (Bin 'k' Empty Empty)
                (Bin 'l' Empty Empty)))
        (Bin 'g' Empty Empty))