breadth first numbering

BFNum.hs

{-# LANGUAGE BangPatterns #-}

-- | Related: <http://www.cs.tufts.edu/~nr/cs257/archive/chris-okasaki/breadth-first.pdf>
module BFNum where

import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap

data Tree a
    = E
    | T a (Tree a) (Tree a)
    deriving Show

-- How do we do a breadth-first search? 
-- The difficult part is rebuilding the tree after we're done.
--
-- So, at the same time we process the tree, we must build a 
-- reconstruction process? Alternatively, we can transpose to
-- some sort of level-tree, then convert back. The latter seems
-- much easier, i.e. we number nodes: 0 for root, 1 for left,
-- 2 for right, 3 for left-left, 4 for left-right, 5 for right-left,
-- 6 for right-right, (n*2+1) and (n*2+2) for children of node n...
type LTree = IntMap

toLTree :: Tree a -> LTree a
toLTree t = toLTree' 0 t IntMap.empty

toLTree' :: Int -> Tree a -> LTree a -> LTree a
toLTree' _ E = id
toLTree' n (T a l r) =
    toLTree' (lN n) l . 
    toLTree' (rN n) r . 
    IntMap.insert n a 

lN, rN :: Int -> Int
lN n = (n * 2) + 1
rN n = (n * 2) + 2

toTree :: LTree a -> Tree a
toTree = toTree' 0 

toTree' :: Int -> LTree a -> Tree a
toTree' n lt = case IntMap.lookup n lt of
    Just a -> T a (toTree' (lN n) lt) (toTree' (rN n) lt)
    Nothing -> E

bfnum :: Tree a -> Tree Int
bfnum = toTree . snd . IntMap.mapAccum fn 0 . toLTree where
    fn !n _ = let n' = n + 1 in (n',n')