Traverse a multiway tree

January 5, 2022

I’m trying to traverse a multiway tree, and to map it’s values like List.map would.

Here is my attempt

type 'a tree = Node of 'a * ('a tree list);;

let tree_map f t = 
  let rec aux acc tr =
    match tr with
    | Node(v, []) -> Node(f v, []) :: acc
    | Node(v, sub) -> let r = List.fold_left aux acc sub in Node(f v, r) :: acc
  in aux [] t
;;

let t = Node (1, [Node (2, [Node (1, [])]); 
                  Node (3, []);
                  Node (1, [Node (5, []); 
                            Node (2, [])])]);;
;;

let res = tree_map (succ) t;;

I tried the suggestion from this answer. Could not implement it successfully.

My function returns this which is obviously incorrect

val res : int tree list =
  [Node (2,
    [Node (2,
      [Node (3, []); Node (6, []); Node (4, []); Node (3, [Node (2, [])])]);
     Node (4, []); Node (3, [Node (2, [])])])]

What’s the problem ?

>Solution :

The answer that you’re referring was about the fold iterator, i.e., when you apply some function to each node of a tree to build some value. The map iterator is a little bit different and is even easier to implement. For each node, you shall return the mapped node.

It is important to keep in mind that the map iterator shall preserve the structure of your data, e.g., when you have the leaf node you need to return the mapped leaf node, e.g.,

| Node (v, []) -> Node (f v, [])

the same is true to the general case, i.e., instead of using the fold iterator, you need to use the map iterator, e.g.,

| Node(v, sub) -> Node (f v, List.map (aux f) sub)

In fact, you don’t even need the empty node case as the general case will work fine for all nodes. Which will make you tree_map a trivial one-liner 🙂

And do not forget to remove the acc parameter from aux, you do not need it.