-
Notifications
You must be signed in to change notification settings - Fork 1
/
LCA.ml
73 lines (59 loc) · 1.83 KB
/
LCA.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
(** Lowest Common Ancestor
By reduction to Range Minimum Query.
We traverse the tree with a DFS and we record the sequence
of depth/node (array [depth] below).
A table [index] maps each node to its first occurrence in
array [depth].
The LCA of [x] and [y] is then the element with minimum
depth in array [depth] between [index[x]] and [index[y]].
Example:
A size = 9
/ \
F E index: A:0 B:2 C:14 D:4 E:13 F:1 G:7 H:10 I:5
/ | \ |
B D H C 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
/ \ depth: 0A 1F 2B 1F 2D 3I 2D 3G 2D 1F 2H 1F 0A 1E 2C 1E 0A
I G
*)
module type TREE = sig
type node
val subtrees: node -> node list
val equal: node -> node -> bool
val hash : node -> int
end
module Make(T: TREE) = struct
module H = Hashtbl.Make(struct
type t = T.node
let equal = T.equal
let hash = T.hash
end)
module E = struct
type t = int * T.node
let compare (d1,_) (d2,_) = Int.compare d1 d2
end
module R = RMQ.Make0(E)
type t = {
index: int H.t;
rmq: R.t;
}
let rec size t =
List.fold_left (fun s t -> s + size t) 1 (T.subtrees t)
let create root =
let n = size root in
let index = H.create n in
let depth = Array.make (2*n - 1) (-1, root) in
let add i p = depth.(i) <- p; i+1 in
let rec fill i ((d, t) as p) =
H.add index t i;
let i = add i p in
let child i c = let i = fill i (d+1, c) in add i p in
List.fold_left child i (T.subtrees t) in
assert (fill 0 (0, root) = 2*n-1);
let rmq = R.create depth in
{ index; rmq }
let lca t x y =
let find n = try H.find t.index n with Not_found -> invalid_arg "lca" in
let ix = find x and iy = find y in
let _, n = R.rmq t.rmq ~lo:(min ix iy) ~hi:(max ix iy + 1) in
n
end