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

A solution to this type of n-ary tree problem is to maintain a list of nodes to process as an argument to the function, along with an accumulator and populate or consume that list of nodes as you go. That way you're only ever worrying about processing one node at a time.

This also has the plus of being tail-recursive.

let rec split tree to_process (acca, accb as acc) =
  match tree, to_process with
  | Empty,              []    -> List.(rev acca, rev accb)
  | Empty,              x::xs -> split x xs acc
  | Vertexa (v, []),    []    -> List.(rev (v::acca), rev accb)
  | Vertexb (v, []),    []    -> List.(rev acca, rev (v::accb))
  | Vertexa (v, y::ys), []    -> split y ys (v::acca, accb)
  | Vertexb (v, y::ys), []    -> split y ys (acca, v::accb) 
  | Vertexa (v, lst),   x::xs -> split x (xs @ lst) (v::acca, accb)
  | Vertexb (v, lst),   x::xs -> split x (xs @ lst) (acca, v::accb)

# split (Vertexa (4, [Vertexb (3, [Vertexa (5, [Empty; Empty])]); Vertexa (1, [])])) [] ([], []);;
- : int list * int list = ([4; 1; 5], [3])

split (Vertexa (4, [Vertexb (3, [Vertexa (5, [Empty; Empty])]); Vertexa (1, [])])) [] ([], [])
split (Vertexb (3, [Vertexa (5, [Empty; Empty])])) [Vertexa (1, [])] (4::[], []) 
split (Vertexa (1, [])) [Vertexa (5, [Empty; Empty])] (4::[], 3::[])
split (Vertexa (5, [Empty; Empty]) [] (1::4::[], 3::[])
split Empty [Empty] (5::1::4::[], 3::[])
split Empty [] (5::1::4::[], 3::[])
([4; 1; 5]; [3])

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 e.g. 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.

Answer 3

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 (c.f., 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.