Source file bt2.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
129
130
131
132
133
134
135
136
137
138
139
140
type (+'a, +'b) t =
  | Leaf
  | Node of int * ('a, 'b) t * 'a * 'b * ('a, 'b) t

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

(** {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 x0 x1 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));
      Node (sl + 1 + sr, l, x0, x1, r)
*)

(** Construct Node *)
let node_ sl l x0 x1 sr r = Node (sl + 1 + sr, l, x0, x1, r)

(** Rotations *)
let rot_left l x0 x1 r k = match r with
  | Node (_, rl, y0, y1, rr) ->
    k (k l x0 x1 rl) y0 y1 rr
  | _ -> assert false

let rot_right l y0 y1 r k = match l with
  | Node (_, ll, x0, x1, lr) ->
    k ll x0 x1 (k lr y0 y1 r)
  | _ -> assert false

(** Balancing *)

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

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

(** Balance a node in both directions: a rotation that fixes one side can
    leave the other side too heavy, so the nodes it builds are balanced the
    same way. *)
let rec node_balance l x0 x1 r =
  let sl = size l and sr = size r in
  (* Fast path: [disbalanced a b] requires [a < b / 2]. *)
  if sr <= 2 * sl + 1 && sl <= 2 * sr + 1 then
    node_ sl l x0 x1 sr r
  else if disbalanced sl sr then
    inc_left l x0 x1 r node_balance
  else if disbalanced sr sl then
    inc_right l x0 x1 r node_balance
  else
    node_ sl l x0 x1 sr r

(** [link l x r] joins balanced trees [l] and [r] of any sizes around [x],
    descending the spine of the larger one until the sizes are compatible.
    Cost is O(log (max sl sr / min sl sr)). *)
let rec link l x0 x1 r =
  let sl = size l and sr = size r in
  (* Fast path: [disbalanced a b] requires [a < b / 2]. *)
  if sr <= 2 * sl + 1 && sl <= 2 * sr + 1 then
    node_ sl l x0 x1 sr r
  else if disbalanced sl sr then
    match r with
    | Node (_, rl, y0, y1, rr) -> node_balance (link l x0 x1 rl) y0 y1 rr
    | Leaf -> assert false
  else if disbalanced sr sl then
    match l with
    | Node (_, ll, y0, y1, lr) -> node_balance ll y0 y1 (link lr x0 x1 r)
    | Leaf -> assert false
  else
    node_ sl l x0 x1 sr r

(** Remove the last element, returned with the remaining tree. *)
let rec pop_max = function
  | Node (_, l, x0, x1, Leaf) -> l, (x0, x1)
  | Node (_, l, x0, x1, r) -> let r, m = pop_max r in node_balance l x0 x1 r, m
  | Leaf -> assert false

(** Remove the first element, returned with the remaining tree. *)
let rec pop_min = function
  | Node (_, Leaf, x0, x1, r) -> (x0, x1), r
  | Node (_, l, x0, x1, r) -> let m, l = pop_min l in m, node_balance l x0 x1 r
  | Leaf -> assert false

(** Public interface *)

let leaf = Leaf

let node l x0 x1 r = link l x0 x1 r

(* The root of the result is the last element of [l] or the first of [r],
   taken from the smaller tree. *)
let join l r = match l, r with
  | Leaf, t | t, Leaf -> t
  | Node (sl, _, _, _, _), Node (sr, _, _, _, _) ->
    if sl <= sr then
      let l, (x0, x1) = pop_max l in link l x0 x1 r
    else
      let (x0, x1), r = pop_min r in link l x0 x1 r

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