OCaml binary tree depth with no stack overflow

Question

I have the following implementation for a binary tree and a depth function to calculate its depth:

type 'a btree =
| Empty
| Node of 'a * 'a btree * 'a btree;;

let rec depth t = match t with
| Empty -> 0
| Node (_, t1, t2) -> 1 + Int.max (depth t1)  (depth t2)

The problem here is that "depth" is recursive and can cause a stack overflow when the tree is too big.

I read about tail recursion and how it can be optimised into a while loop by the compiler to remove the stack call.

How would you make this function tail recursive or make it use a while/for loop instead?

Answer 1

type 'a btree =
| Empty
| Node of 'a * 'a btree * 'a btree;;

let max x y = if x > y then x else y

let depth t =
  let rec dep m = function (* d records current level, m records max depth so far *)
    | [] -> m
    | (Empty,d)::tl -> dep (max m d) tl
    | (Node (_,l,r),d)::tl -> dep (max m d) ((l,d+1)::(r,d+1)::tl)
  in 
  dep 0 [(t,0)]

Basically, you need 3 things:

1. a list (stack) to store nodes along the paths
2. a indicator to record the current depth
3. the current max depth so far

Whenever we face a problem that needs to remove the possible stackoverflow problem, we should think two things: tail-recursive and explicit stack.

For tail-recursive, you have to find a way to explicitly store the temporary data generated through each recursion step.

For explicit stack, remember the reason that recursion can work is because internally it uses a stack with a limited size. If we analyse the logic and make that stack explicit, we then don't need that internal stack any more.

Answer 2

In pragmatic cases the solution is to use a balanced tree, which limits the depth to some multiple of log(n). Even for very large n, log(n) is small enough that you won't run out of stack space.

Otherwise see the SO page linked by Kadaku. It has ridiculously good answers to the question.

Answer 3

I already answered similar question once. Reposting the solution:

There's a neat and generic solution using fold_tree and CPS - continuous passing style:

let fold_tree tree f acc =
  let loop t cont =
    match tree with
    | Leaf -> cont acc
    | Node (x, left, right) ->
      loop left (fun lacc ->
        loop right (fun racc ->
          cont @@ f x lacc racc))
  in loop tree (fun x -> x)

let depth tree = fold_tree tree (fun x dl dr -> 1 + (max dl dr)) 0