How do I sum all the elements from the leaf to the root of each branch in a n-ary tree in OCaml?

Question

I'm trying to make a function in OCaml where, given a n-tree, it returns a list of all the sums from leaf to root of all the branches. This is what i did:

exception NotFound

type 'a ntree = Tr of 'a * 'a ntree list

let leaf x = Tr (x, [])

let alb = 
  Tr (1, [Tr (2, [leaf 3; leaf 4; leaf 2]); 
          Tr (5, [leaf 11; leaf 10]); 
          Tr (3, [leaf 9; leaf 7; leaf 10])])

let rec peso (Tr (x, tlist)) =
  match tlist with
    [] -> [x]
  | _ -> [x + peso_l tlist]
and peso_l = function 
    [] -> raise NotFound
  | [t] -> peso t
  | t::rest -> peso t :: peso_l rest

But it doesn't work because, I think,

| _ ->[x + peso_l tlist]

returns something like [x+ [t]] (am I right?). How can I fix it?

Answer 1

When you write [x + peso_l tlist] what you want to do is add x to each element of the list returned by peso_l tlist. This can be achieved with List.map:

exception NotFound

type 'a ntree = Tr of 'a * 'a ntree list

let leaf x = Tr (x, [])

let alb =
  Tr
    ( 1,
      [
        Tr (2, [ leaf 3; leaf 4; leaf 2 ]);
        Tr (5, [ leaf 11; leaf 10 ]);
        Tr (3, [ leaf 9; leaf 7; leaf 10 ]);
      ] )

let rec peso (Tr (x, tlist)) =
  match tlist with [] -> [ x ] | _ -> List.map (( + ) x) (peso_l tlist)
and peso_l = function
  | [] -> raise NotFound
  | [ t ] -> peso t
  | t :: rest -> peso t @ peso_l rest

let () =
  Format.printf "@[<v 0>%a@."
    Format.(
      pp_print_list ~pp_sep:pp_print_cut (fun ppf d ->
          Format.fprintf ppf "%d" d))
    (peso alb)