-
Notifications
You must be signed in to change notification settings - Fork 20
/
mo-tree.lisp
183 lines (175 loc) · 7.79 KB
/
mo-tree.lisp
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
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
(defpackage :cp/mo-tree
(:use :cl)
(:export #:mo-tree #:mo-tree-p #:make-mo-tree #:mo-tree-get-lca
#:mo-tree-get-current #:mo-tree-process2)
(:documentation "Provides Mo's algorithm for paths on a tree with weighted
vertex.
Reference:
https://ei1333.hateblo.jp/entry/2017/09/11/211011 (Japanese)"))
(in-package :cp/mo-tree)
;; not tested
(defstruct (mo-tree (:constructor %make-mo-tree)
(:conc-name %mo-tree-))
;; LCA part
(graph nil :type (simple-array list (*)))
(max-level nil :type (integer 0 #.most-positive-fixnum))
(depths nil :type (simple-array fixnum (*)))
(parents nil :type (simple-array fixnum (* *)))
(euler-tour nil :type (simple-array (integer 0 #.most-positive-fixnum) (*)))
;; Mo part
(flags nil :type simple-bit-vector)
(lefts nil :type (simple-array (integer 0 #.most-positive-fixnum) (*)))
(rights nil :type (simple-array (integer 0 #.most-positive-fixnum) (*)))
(orders nil :type (simple-array (integer 0 #.most-positive-fixnum) (*)))
(lcas nil :type (simple-array (integer 0 #.most-positive-fixnum) (*)))
(index 0 :type (integer 0 #.most-positive-fixnum))
(posl 0 :type (integer 0 #.most-positive-fixnum))
(posr 0 :type (integer 0 #.most-positive-fixnum)))
(defun mo-tree-get-lca (mo vertex1 vertex2)
"Returns the lowest common ancestor of the two vertices."
(declare (optimize (speed 3))
((integer 0 #.most-positive-fixnum) vertex1 vertex2))
(let* ((u vertex1)
(v vertex2)
(depths (%mo-tree-depths mo))
(parents (%mo-tree-parents mo))
(max-level (%mo-tree-max-level mo)))
(declare (fixnum u v))
;; Ensure depth[u] <= depth[v]
(when (> (aref depths u) (aref depths v))
(rotatef u v))
(dotimes (k max-level)
(when (logbitp k (- (aref depths v) (aref depths u)))
(setq v (aref parents v k))))
(if (= u v)
u
(loop for k from (- max-level 1) downto 0
unless (= (aref parents u k) (aref parents v k))
do (setq u (aref parents u k)
v (aref parents v k))
finally (return (aref parents u 0))))))
(defun make-mo-tree (graph bucket-width lefts rights)
"LEFTS := vector of vertices of an end of queries (inclusive)
RIGHTS := vector of vertices of another end of queries (inclusive)
BUCKET-WIDTH would be better set to 2N/sqrt(Q) where N is the width of the
universe and Q is the number of queries."
(declare (optimize (speed 3))
(vector lefts rights)
((integer 0 #.most-positive-fixnum) bucket-width)
(inline sort))
(let* ((graph (coerce graph '(simple-array list (*))))
(n (length graph))
(max-level (+ 1 (integer-length (- n 2))))
(depths (make-array n :element-type 'fixnum :initial-element -1))
(parents (make-array (list n max-level) :element-type 'fixnum))
(euler-tour (make-array (- (* 2 n) 1) :element-type '(integer 0 #.most-positive-fixnum)))
(ins (make-array n :element-type '(integer 0 #.most-positive-fixnum)))
(end 0))
(declare ((integer 0 #.most-positive-fixnum) n max-level end))
;; LCA part
(labels ((dfs (v parent depth)
(declare (fixnum v parent))
(setf (aref ins v) end
(aref depths v) depth
(aref parents v 0) parent
(aref euler-tour end) v
end (+ end 1))
(dolist (child (aref graph v))
(unless (= child parent)
(dfs child v (+ depth 1))
(setf (aref euler-tour end) child
end (+ end 1))))))
(dfs 0 -1 0)
(dotimes (k (- max-level 1))
(dotimes (v n)
(setf (aref parents v (+ k 1))
(if (= -1 (aref parents v k))
-1
(aref parents (aref parents v k) k)))))
;; Mo part
(let* ((q (length lefts))
(res-lefts (make-array q :element-type '(integer 0 #.most-positive-fixnum)))
(res-rights (make-array q :element-type '(integer 0 #.most-positive-fixnum)))
(orders (make-array q :element-type '(integer 0 #.most-positive-fixnum)))
(lcas (make-array q :element-type '(integer 0 #.most-positive-fixnum)))
(mo (%make-mo-tree :graph graph
:max-level max-level
:depths depths
:parents parents
:euler-tour euler-tour
:flags (make-array n :element-type 'bit :initial-element 0)
:lefts res-lefts
:rights res-rights
:orders orders
:lcas lcas)))
(assert (= q (length rights)))
(dotimes (i q)
(let ((l (aref lefts i))
(r (aref rights i)))
(declare ((integer 0 #.most-positive-fixnum) l r))
(when (> l r) (rotatef l r))
(setf (aref res-lefts i) (+ 1 (aref ins l))
(aref res-rights i) (+ 1 (aref ins r))
(aref lcas i) (mo-tree-get-lca mo l r))))
(dotimes (i q)
(setf (aref orders i) i))
(setf (%mo-tree-orders mo)
(sort orders
(lambda (x y)
(if (= (floor (aref res-lefts x) bucket-width)
(floor (aref res-lefts y) bucket-width))
;; Even number [Odd number] block is in ascending
;; [desceding] order w.r.t. the right end.
(if (evenp (floor (aref res-lefts x) bucket-width))
(< (aref res-rights x) (aref res-rights y))
(> (aref res-rights x) (aref res-rights y)))
(< (aref res-lefts x) (aref res-lefts y))))))
mo))))
(declaim (inline mo-tree-get-current))
(defun mo-tree-get-current (mo)
"Returns the original index of the current (not yet processed) query."
(aref (%mo-tree-orders mo) (%mo-tree-index mo)))
(declaim (inline mo-tree-process2))
(defun mo-tree-process2 (mo extend shrink)
"Processes the next query. EXTEND and SHRINK take an argument: the index to be
added or removed."
(declare (function extend shrink))
(let* ((orders (%mo-tree-orders mo))
(index (%mo-tree-index mo))
(ord (aref orders index))
(lcas (%mo-tree-lcas mo))
(left (aref (%mo-tree-lefts mo) ord))
(right (aref (%mo-tree-rights mo) ord))
(posl (%mo-tree-posl mo))
(posr (%mo-tree-posr mo))
(flags (%mo-tree-flags mo))
(euler-tour (%mo-tree-euler-tour mo)))
(declare ((integer 0 #.most-positive-fixnum) posl posr))
(labels ((frob (vertex)
(declare ((integer 0 #.most-positive-fixnum) vertex))
(if (zerop (aref flags vertex))
(funcall extend vertex)
(funcall shrink vertex))
(setf (aref flags vertex)
(logxor 1 (aref flags vertex)))))
;; remove LCA of the previous query
(unless (zerop index)
(frob (aref lcas (aref orders (- index 1)))))
(loop while (< left posl)
do (decf posl)
(frob (aref euler-tour posl)))
(loop while (< posr right)
do (frob (aref euler-tour posr))
(incf posr))
(loop while (< posl left)
do (frob (aref euler-tour posl))
(incf posl))
(loop while (< right posr)
do (decf posr)
(frob (aref euler-tour posr)))
;; process LCA separately
(frob (aref lcas (aref orders index))))
(setf (%mo-tree-posl mo) posl
(%mo-tree-posr mo) posr)
(incf (%mo-tree-index mo))
ord))