Binary Tree Depth

Tree.purs

module BinaryTree where

import Prelude

import Data.Maybe (Maybe(..))
import Data.Tuple (Tuple(..))

data Tree a
  = Empty
  | Node (Tree a) a (Tree a)

foldl :: ∀ a b. (b -> a -> b) -> b -> Tree a -> b
foldl f init Empty = init
foldl f init (Node l x r) = f (foldl f (foldl f init l) r) x

map :: ∀ a b. (a -> b) -> Tree a -> Tree b
map f Empty = Empty
map f (Node l x r) = Node (map f l) (f x) (map f r)

mapWithDepth :: ∀ a b. (Int -> a -> b) -> Tree a -> Tree b
mapWithDepth f tree =
  go f 0 tree
  where
    go f depth Empty = Empty
    go f depth (Node l x r) =
      Node (go f (depth + 1) l) (f depth x) (go f (depth + 1) r)

zipWithDepth :: ∀ a. Tree a -> Tree { val :: a, depth :: Int }
zipWithDepth = mapWithDepth (\d v -> { val : v, depth : d})

maxDepth :: ∀ a. Tree a -> Maybe { val :: a, depth :: Int }
maxDepth tree =
  foldl f Nothing $ zipWithDepth tree
  where
    f Nothing curr = curr
    f (Just acc) curr =
      if acc.depth >= curr.depth then acc
      else curr


fastMaxDepth :: ∀ a. Tree a -> Maybe { val :: a, depth :: Int }
fastMaxDepth tree =
  go 1 tree
  where
    go depth Empty = Nothing
    go depth (Node l x r) =
      let
        ld = go (depth + 1) l
        rd = go (depth + 1) r
      in case (Tuple ld rd) of
        (Tuple Nothing Nothing) -> Just { val : x, depth : depth}
        (Tuple l Nothing) -> l
        (Tuple Nothing r) -> r
        (Tuple (Just lval) (Just rval)) ->
          if lval.depth >= rval.depth then Just lval
          else Just rval