let rec updateTree index newElement tree size = match tree with
  | Leaf element -> Leaf newElement
  | Node {element; left; right} ->
    let newSize = ((size + 1) / 2) - 1 in
    if index == 0 then Node {element = newElement; left; right}
    else if index <= newSize then
      Node {
        element;
        left = updateTree (index - 1) newElement left newSize;
        right;
      }
    else
      Node {
        element;
        left;
        right = updateTree (index - newSize - 1) newElement right newSize;
      }
;;

let rec update index element ls = match ls with
  | [] -> raise IndexOutOfBoundsException
  | {size; tree} :: rest -> if index < size
    then {size; tree = updateTree index element tree size} :: rest
    else {size; tree} :: (update (index - size) element rest)
;;