Source file mbt.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
module type MEASURE = sig
  type +'a measurable
  type measure
  val empty : measure
  val cat : measure -> 'a measurable -> measure -> measure
end

module Make(M : MEASURE) = struct

  type 'a t =
    | Leaf
    | Node of int * 'a t * 'a M.measurable * 'a t * M.measure

  let size = function
    | Node (s, _, _, _, _) -> s
    | Leaf -> 0

  let measure = function
    | Node (_, _, _, _, m) -> m
    | Leaf -> M.empty

  (** {1 Balance criteria}
      Functions are not symmetric.
      The first argument should always be of the same power of two or smaller
      (guaranteed by construction). *)

  (** [smaller_ell smin smax] iff
      - [smin] is less than [smax]
      - [smin] and [smax] differs by less than two magnitude orders, i.e
        msbs(smin) >= msbs(smax) - 1
      where msbs is the index of the most significant bit set *)
  let smaller_ell smin smax = (smin < smax) && ((smin land smax) lsl 1 < smax)

  (** [disbalanced smin smax] check if two sub-trees of size [smin] and [smax],
      are disbalanced. That is, msbs(smin) < msbs(smax) - 1 *)
  let disbalanced smin smax = smaller_ell smin (smax lsr 1)

  (** {1 Smart but not too much constructors} *)

  (** Construct node and check balance
      let node_ l x r =
        let sl = size l and sr = size r in
        if sl < sr then
          assert (not (disbalanced sl sr))
        else
          assert (not (disbalanced sr sl));
        let ml = measure l and mr = measure r in
        Node (sl + 1 + sr, l, x, r, M.cat ml x mr)
  *)

  (** Construct Node *)
  let node_ l x r =
    Node (size l + 1 + size r, l, x, r, M.cat (measure l) x (measure r))

  (** Rotations *)
  let rot_left l x r k = match r with
    | Node (_, rl, y, rr, _) ->
      k (k l x rl) y rr
    | _ -> assert false

  let rot_right l y r k = match l with
    | Node (_, ll, x, lr, _) ->
      k ll x (k lr y r)
    | _ -> assert false

  (** Balancing *)

  let inc_left l x r k =
    let r = match r with
      | Node (_, rl, y, rr, _) when smaller_ell (size rr) (size rl) ->
        rot_right rl y rr k
      | _ -> r
    in
    rot_left l x r k

  let inc_right l y r k =
    let l = match l with
      | Node (_, ll, x, lr, _) when smaller_ell (size ll) (size lr) ->
        rot_left ll x lr k
      | _ -> l
    in
    rot_right l y r k

  (** Balance trees leaning to the right *)
  let rec node_left l x r =
    if disbalanced (size l) (size r) then
      inc_left l x r node_left
    else
      node_ l x r

  (** Balance trees leaning to the left *)
  let rec node_right l y r =
    if disbalanced (size r) (size l) then
      inc_right l y r node_right
    else
      node_ l y r

  (** Public interface *)

  let leaf = Leaf

  let node l x r = match l, r with
    | Leaf, Leaf -> node_ leaf x leaf
    | l, r when size l < size r ->
      node_left l x r
    | l, r ->
      node_right l x r

  let rec join l r = match l, r with
    | Leaf, t | t, Leaf -> t
    | Node (sl, ll, x, lr, _), Node (sr, rl, y, rr, _) ->
      if sl <= sr then
        node (join l rl) y rr
      else
        node ll x (join lr r)

  let rec rank n = function
    | Leaf -> raise Not_found
    | Node (_, l, x, r, _) ->
      let sl = size l in
      if n = sl then
        x
      else if n < sl then
        rank n l
      else
        rank (n - 1 - sl) r

end