(* persistent implementation using a map giving the class id 
 * the updating of the map is implemented with a customized remap function
 * which preserves the AVL structure
 *)

module M = struct

  type t =
      Empty
    | Node of t * int * int * t * int
	
  let height = function
      Empty -> 0
    | Node(_,_,_,_,h) -> h

  let create l x d r =
    let hl = height l and hr = height r in
    Node(l, x, d, r, (if hl >= hr then hl + 1 else hr + 1))

  let bal l x d r =
    let hl = match l with Empty -> 0 | Node(_,_,_,_,h) -> h in
    let hr = match r with Empty -> 0 | Node(_,_,_,_,h) -> h in
    if hl > hr + 2 then begin
      match l with
          Empty -> invalid_arg "Map.bal"
        | Node(ll, lv, ld, lr, _) ->
            if height ll >= height lr then
              create ll lv ld (create lr x d r)
            else begin
              match lr with
                  Empty -> invalid_arg "Map.bal"
		| Node(lrl, lrv, lrd, lrr, _)->
                    create (create ll lv ld lrl) lrv lrd (create lrr x d r)
            end
    end else if hr > hl + 2 then begin
      match r with
          Empty -> invalid_arg "Map.bal"
        | Node(rl, rv, rd, rr, _) ->
            if height rr >= height rl then
              create (create l x d rl) rv rd rr
            else begin
              match rl with
                  Empty -> invalid_arg "Map.bal"
		| Node(rll, rlv, rld, rlr, _) ->
                    create (create l x d rll) rlv rld (create rlr rv rd rr)
            end
    end else
      Node(l, x, d, r, (if hl >= hr then hl + 1 else hr + 1))

  let empty = Empty

  let rec add x data = function
      Empty ->
        Node(Empty, x, data, Empty, 1)
    | Node(l, v, d, r, h) ->
        let c = compare x v in
        if c = 0 then
          Node(l, x, data, r, h)
        else if c < 0 then
          bal (add x data l) v d r
        else
          bal l v d (add x data r)
	    
  let rec find x = function
      Empty ->
        raise Not_found
    | Node(l, v, d, r, _) ->
        let c = compare x v in
        if c = 0 then d
        else find x (if c < 0 then l else r)

  let rec remap d1 d2 = function
    | Empty -> Empty
    | Node (l, v, d, r, h) -> 
	let d = if d == d1 then d2 else d in
	Node (remap d1 d2 l, v, d, remap d1 d2 r, h)

end

type t = M.t

let create n = M.empty

let find m i = try M.find i m with Not_found -> i
    
let union m i j =
  let ri = find m i in
  let rj = find m j in
  if ri = rj then
    m
  else
    M.remap ri rj (M.add ri rj m)