Balanced binary (AA) search trees

aa.tq

type rec AATree a = E | T(Number, AATree a, a, AATree a)

let empty = E

let isEmpty = [E → True | T _ → False]

let rec count =
  [ E → 0
  | T(_, l, _, r) → count l + 1 + count r ]

let rec contains k =
  [ E → False
  | T(_, l, v, r) →
      compare(k, v)
      @ [ Less → contains k l
        | Equal → True
        | Greater → contains k r ] ]

let skew =
  [ T(n, T(ln, ll, lv, lr), v, r) as t →
      if n=ln then T(ln, ll, lv, T(n, lr, v, r)) else t
  | t → t ]

let lvl = [E → 0 | T(n, _, _, _) → n]

let split =
  [ T(n, l, v, T(rn, rl, rv, rr)) as t →
      if n=lvl rr then T(rn+1, T(n, l, v, rl), rv, rr) else t
  | t → t ]

let left f =
  [ T(n, l, v, r) → T(n, f l, v, r)
  | E → E ]

let right f =
  [ T(n, l, v, r) → T(n, l, v, f r)
  | E → E ]

let rec add x =
  [ E → T(1, E, x, E)
  | T(n, l, v, r) as t →
      compare(x, v)
      @ [ Less → split(skew(left (add x) t))
        | Equal → T(n, l, x, r)
        | Greater → split(skew(right (add x) t)) ] ]

let fixup(n, l, v, r) =
  T(n, l, v, r)
  @ skew
  @ right skew
  @ right (right skew)
  @ split
  @ right split

let rec first =
  [ E → None
  | T(_, l, v, _) → first l @ Option.default v @ Some ]

let rec last =
  [ E → None
  | T(_, _, v, r) → last r @ Option.default v @ Some ]

let rec remove x =
  [ E → E
  | T(n, l, v, r) as t →
      compare(x, v)
      @ [ Less → left (remove x) t
        | Equal → merge n (l, r)
        | Greater → right (remove x) t ]
      @ [ T(n, l, v, E) as t →
            if lvl l < n-1 || 0 < n-1 then fixup(n-1, l, v, E) else t
        | T(n, l, v, T(rn, rl, rv, rr)) as t →
            if lvl l < n-1 || rn < n-1 then
              let n = n-1 in
              let rn = if rn > n then n else rn in
              fixup(n, l, v, T(rn, rl, rv, rr))
            else t
        | t → t ] ]
and merge n (l, r) =
  last l
  @ [ None →
        first r
        @ [ None → E
          | Some v → T(n, l, v, remove v r) ]
    | Some v → T(n, remove v l, v, r) ]

let n = 500

let xs = Array.init n [_ → Random.next n]

let aa = time "add" (Array.fold [s, x → add x s] empty) xs

let () = example("count aa", count aa)

let set = time "Set.add" (Array.fold [s, x → Set.add x s] (Set.empty())) xs

let () = example("count set", Set.count set)


let aa =
  let aa = for 1 1 n empty [s, x → add x s] in
  time "remove" (Array.fold [s, x → remove x s] aa) xs

let () = example("count aa", count aa)

let set =
  let set = for 1 1 n (Set.empty()) [s, x → Set.add x s] in
  time "Set.remove" (Array.fold [s, x → Set.remove x s] set) xs

let () = example("count set", Set.count set)