basic treap implementation

treap.ml

 (** Task X: Treap. *)

module Treap = struct
  type ('a, 'b) t =
    | Empty
    | Node of ('a, 'b) t * 'a * 'b * ('a, 'b) t

  let empty = Empty

  let rec merge l r = match (l, r) with
    | (Empty, _) -> r
    | (_, Empty) -> l
    | (Node (l1, key1, rank1, r1),
       Node (l2, key2, rank2, r2)) ->
      if rank1 < rank2 then
        Node (l1, key1, rank1, merge r1 r)
      else
        Node (merge l l2, key2, rank2, r2)

  let rec split t split_key = match t with
    | Empty -> (Empty, Empty)
    | Node (l, key, rank, r) ->
      if key > split_key then
        let (ls1, ls2) = split l split_key in
        (ls1, Node (ls2, key, rank, r))
      else
        let (rs1, rs2) = split r split_key in
        (Node (l, key, rank, rs1), rs2)

  let add t key rank =
    let (l, r) = split t key in
    let mid = Node (Empty, key, rank, Empty) in
    merge (merge l mid) r
end


let () =
  let n = Scanf.scanf "%i " (fun x -> x) in
  let map = Hashtbl.create n in

  let int_of_node = function
    | Treap.Empty -> 0
    | Treap.Node (_, key, _, _) -> Hashtbl.find map key
  in

  let rec traverse parent aux = function
    | Treap.Empty -> ()
    | (Treap.Node (l, key, _, r)) as node -> begin
      traverse node aux l;
      aux.(int_of_node node) <- ((int_of_node parent),
                                 (int_of_node l),
                                 (int_of_node r));
      traverse node aux r
    end
  in

  let build t =
    let rec inner t = function
      | i when i = n -> t
      | i ->
        let (key, rank) = Scanf.scanf "%i %i " (fun x y -> (x, y)) in
        Hashtbl.add map key (i + 1);
        inner (Treap.add t key rank) (i + 1)
    in inner t 0
  in

  let t = build Treap.empty in
  let aux = Array.make (n + 1) (0, 0, 0) in
  begin
    traverse Treap.Empty aux t;

    print_endline "YES";
    for i = 1 to n do
      let (p, l, r) = aux.(i) in
      Printf.printf "%i %i %i\n" p l r
    done
  end