let trap heights =
  let n = Array.length heights in
  if n < 3 then 0
  else begin
    let left_max = Array.make n 0 in
    let right_max = Array.make n 0 in
    left_max.(0) <- heights.(0);
    for i = 1 to n - 1 do
      left_max.(i) <- if heights.(i) > left_max.(i - 1) then heights.(i) else left_max.(i - 1)
    done;
    right_max.(n - 1) <- heights.(n - 1);
    for i = n - 2 downto 0 do
      right_max.(i) <- if heights.(i) > right_max.(i + 1) then heights.(i) else right_max.(i + 1)
    done;
    let water = ref 0 in
    for i = 0 to n - 1 do
      let min_lr = if left_max.(i) < right_max.(i) then left_max.(i) else right_max.(i) in
      water := !water + min_lr - heights.(i)
    done;
    !water
  end

;;

trap [| 0; 1; 0; 2; 1; 0; 1; 3; 2; 1; 2; 1 |]