How do I replace part of a tree with another tree at the specified index in OCaml?

Question

I have the following tree structure:

type 'a tree =
    | Function of string * 'a tree list  (* function name and arguments *)
    | Terminal of 'a

I use this tree structure to construct an abstract syntax tree:

type atoms =
    | Int of int
    | Bool of bool

let t1 = Function ("+", [Function ("*", [Terminal (Int 5);
                                         Terminal (Int 6)]);
                         Function ("sqrt", [Terminal (Int 3)])])

let t2 = Function ("-", [Terminal (Int 4); Terminal (Int 2)])

Tree representation of t1:

Tree representation of t2:

Goal: replace one of the subtrees in t1 with t2 at a specified t1 index position. The index position starts at 0 at the root node and is depth-first. In the figures above, I have labelled all the nodes with their index to show this.

For example, replace_subtree t1 4 t2 replaces the subtree at index 4 in t1 with t2, resulting in this tree:

Function ("+", [Function ("*", [Terminal (Int 5);
                                Terminal (Int 6)]);
                Function ("-", [Terminal (Int 4);
                                Terminal (Int 2)])])

This is essentially a crossover operation in tree-based genetic programming.

How do I implement replace_subtree in OCaml?

I would strongly prefer a purely functional solution.

---

Note that this question is similar to How do I replace part of a tree with another tree at the specified index?, except that the programming language in this question is OCaml instead of Scheme/Racket. I have some trouble understanding Scheme/Racket, so I am looking for an OCaml solution.

Answer 1

Let's say you had a recursive function dfs that visited every node of a tree, with one parameter being the index number of the node.

Now rewrite this function to return an additional value which is a copy of the subtree below the node. I.e, it visits the subtrees of the node recursively (receiving copies of the subtrees) and constructs a new node as their parent.

Now add two parameters to the function, the index and the desired replacement. When reaching the desired index, the function returns the replacement instead of the copy of the node.

(Since this looks like possible homework I don't want to provide code.)

Answer 2

I have written a solution:

let rec replace_subtree' start_index tree replacement_index replacement
        : (int * 'a tree) =
    (* Returns (max_index, new_tree), where max_index = start_index + number of
       nodes in new_tree - 1, and where the replacement is counted as a single
       node. *)
    if start_index = replacement_index then
        (start_index, replacement)
    else
        match tree with
        | Function (name, args) ->
            (* (start_index + 1) to account for this function node itself. *)
            let (max_index, new_args) = replace_subtree_args (start_index + 1)
                                                             args
                                                             replacement_index
                                                             replacement in
            (max_index, Function (name, new_args))
        | Terminal _ ->
            (start_index, tree)

and replace_subtree_args arg_index args replacement_index replacement
        : (int * 'a tree list) =
    (* `arg_index` is the index of the first item in `args` (note that `args`
       could be empty, however).
       Returns (max_index, replaced_args), where max_index = arg_index +
       number of nodes in all transformed args - 1, and where the replacement is
       counted as a single node. *)
    let rec f arg_index args acc =
        match args with
        | [] -> (arg_index - 1, List.rev acc)
        | arg::rest_args ->
            let (max_index, arg_result) = replace_subtree' arg_index
                                                           arg
                                                           replacement_index
                                                           replacement in
            f (max_index + 1) rest_args (arg_result::acc)
    in
    f arg_index args []

let replace_subtree = replace_subtree' 0

Example usage:

let string_of_terminal = function
    | Int x -> string_of_int x
    | Bool b -> string_of_bool b

let rec string_of_tree = function
    | Function (name, args) ->
        "(" ^
        String.concat " " (name::(List.map string_of_tree args)) ^
        ")"
    | Terminal x -> string_of_terminal x

let () =
    List.iter (fun n ->
                  let (max_index, new_tree) = replace_subtree t1 n t2 in
                  print_string ("Index " ^ (string_of_int n) ^ ":  ");
                  print_endline (string_of_tree new_tree))
              (List.init 8 Fun.id)

Result:

Index 0:  (- 4 2)
Index 1:  (+ (- 4 2) (sqrt 3))
Index 2:  (+ (* (- 4 2) 6) (sqrt 3))
Index 3:  (+ (* 5 (- 4 2)) (sqrt 3))
Index 4:  (+ (* 5 6) (- 4 2))         ; <- Here.
Index 5:  (+ (* 5 6) (sqrt (- 4 2)))
Index 6:  (+ (* 5 6) (sqrt 3))
Index 7:  (+ (* 5 6) (sqrt 3))

Better solutions are most welcome.

Answer 3

Not sure if it's better than the solution you proposed...

I used on less recursive function than you :

let replace to_replace index replacement =
  let rec dfs i tree =
    if i = index then (replacement, i + 1) (* we can replace *)
    else if i > index then (tree, i) (* we already replaced *)
    else
      match tree with
      | Terminal _ -> (tree, i + 1)
      | Function (n, children) ->
          let new_i, new_children = iter_children (i + 1) children in
          (Function (n, new_children), new_i)
  and iter_children i = function
    | [] -> (i, [])
    | child :: children ->
        let new_child, new_i = bfs i child in
        if new_i = index + 1 then (new_i + 1, new_child :: children)
          (* +1 to stop the bfs after appending the children to the Function node *)
        else
          let last_i, last_children = iter_children new_i children in
          (last_i, new_child :: last_children)
  in
  fst @@ bfs 0 to_replace