I'd work with a tree monadized with a Continuation_monad, that is, with:

#use "path/to/tree_monadize.ml";;

You could in theory also use the richer "monads.ml" library:

#use "path/to/monads.ml";;
module TC = Tree_monad.T(Continuation_monad);;
let tree_monadize = TC.distribute;;

But I found it easier to work out the solution using the simpler "tree_monadize.ml" library. The "monadize.ml" library enforces encapsulation in a way that's nice and sanitized, but makes experimentation more cumbersome.

It makes debugging easier if we work with a tree whose starting leaf elements are differently typed than the tree we aim to finish with. So:

let tree = Node(Leaf '1', Node(Leaf '2', Node(Leaf '3', Leaf '1')));;

Now, we already know how to count the leaves using a Continuation monad in tree shape:

let v0 = TreeCont.monadize (fun a k -> 1 + k a) tree (fun t -> 0);;

There are two interesting functions here: we'll call them the "distributed" function fun a k -> 1 + k a and the "seed" function fun t -> 0.

Our first step is to extend v0 so that we count how many leaves there are of each value, rather than how many leaves in total:

let update_env e x = fun y -> (if x = y then 1 else 0) + e y;;

let v1 = TreeCont.monadize (fun a k ->
  fun e0 ->
    let ecur = k a e0
    in update_env ecur a
) tree (fun t e -> e) (fun a -> 0);;

(* now
   v1 '0' ~~> 0
   v1 '1' ~~> 2
   v1 '2' ~~> 1
 *)

How does this work? Our distributed function fun a k -> ... takes a leaf element a and a continuation k, which maps leaf elements and environments e0 to new environments ecur. It gives back a function from e0 to an updated version of ecur.

Our seed function here is the initial continuation. Instead of taking a leaf element and an env, it takes a tree of such elements and an env. In general, wherever the distributed function takes 'as, the seed function takes 'a trees.

Here our seed function just ignores the tree and returns the env unchanged.

What that gives us is a function from environments to environments. We give it an initial environment fun a -> 0 to get started. We get back an environment v1 that maps each leaf element to how many times that leaf element appeared in the tree we started with.

OK? You might want to stop reading here, and see how much further you can get on your own.

But there are more hints if you need them.