OCaml function with data type tree

Question

We are give a tree that contains two types of elements. It is defined with the following data-structure.

type ( 'a , 'b ) tree =
   Empty
   | Vertexa of 'a * ( 'a , 'b ) tree list
   | Vertexb of 'b * ( 'a , 'b ) tree list

Write a function split: ('a,'b) tree -> 'a list * 'b list, that saves all elements of type 'a to the first list and all elements of type 'b in the second list.

I had an idea of doing this recursively but I am kind of stuck on this. I will attach my attempt even though it does not work at all :/

let rec one drevo_list=
    match drevo_list with
   | [Empty]->Empty 
   | Empty::tl -> one tl
   | Vertexa(a,b)::tl -> Vertexa(a,b@tl) 
   | Vertexb(a,b)::tl -> Vertexb(a,b@tl) 

This function turns a list into a tree. I needed it for the recursion, since the second parametar in Vertexa or Vertexb is a list. And this works But the recursion part does not.

let rec split drevo=
   match drevo with
   | Empty -> [],[]
   | Vertexa(a,b)-> split (one b)
   | Vertexb(a,b)-> split (one b)

This part does not work and I have no idea how to finish it. Does any one have an idea how to finish this?

Answer 1

You don't need the drevo_list function to solve this problem. It will actually lead you in a wrong direction.

You need to use List.map to apply your split on a list of trees. You will get a value of ('a list * 'b list) list type. Now you need a helper function concat_pairs that will flatten this value into a pair of type 'a list * 'b list (cf, standard concat function). To implement this function you may use List.fold_left . The rest is trivial.

Note, of course this is a greedy solution. When you finished with it, you may try to find a better solution, that is more efficient and tail recursive.

Answer 2

There are at least two parts of this function that make it difficult to write:

1. A function that returns a pair of lists needs to pack and unpack its return value in each recursive step through eg helper functions, match statements or let bindings. One way would be to write a function that inserts an element into a list inside a pair:

    let insertA a (xs, ys) = (a::xs, ys) let insertB b (xs, ys) = (xs, b::ys) 

2. A function that is recursive over both a tree type and an embedded list type requires the combination of two recursion patterns. This can be solved using either a set of mutually recursive functions, or using higher-order combinators for the lists. Here is an outline of a solution using the former strategy:

    let rec split s = match s with | Empty -> ([], []) | Vertexa (a, ts) -> (* if we had just one t: insertA a (split t) *) | Vertexb (a, ts) -> (* if we had just one t: insertB b (split t) *) 

   So you need a function splitMany : ('a, 'b) tree list -> ('a list, 'b list) that may call back on split for each of its individual trees.

    and rec splitMany ts = match ts with | [] -> ([], []) | (t:ts') -> (* merge (split t) with (splitMany ts') *) 

   For the higher-order function approach, you can avoid the explicit mutual recursion by having the function pass itself to a set of higher-order functions and thus not entangle it in the implementation of the higher-order functions:

    let rec split s = match s with | Empty -> [],[] | Vertexa (a, ts) -> insertA (concat_pairs (map split ts)) | Vertexb (a, ts) -> insertB (concat_pairs (map split ts)) 

   where concat_pairs is ivg's invention.