dkh.segtree.lazyseg source code

1 module dkh.segtree.lazyseg;
2 
3 import dkh.segtree.primitive;
4 
5 import std.functional : binaryFun;
6 
7 /**
8 遅延伝搬SegTree
9 
10 T型の配列aに対して、a[l..r] += x(xはL型)、opTT(a[l..r])が高速に計算できる
11 
12 Params:
13     opTT = (T, T)の演算(結果をまとめる)
14     opTL = (T, L)の演算(クエリを適用する)
15     opLL = (L, L)の演算(クエリをまとめる)
16     eT = Tの単位元
17     eL = Lの単位元
18 */
19 alias LazySeg(T, L, alias opTT, alias opTL, alias opLL, T eT, L eL, alias Engine = LazySegEngine) =
20     SegTree!(Engine, T, L , binaryFun!opTT, binaryFun!opTL, binaryFun!opLL, eT, eL);
21 
22 ///
23 unittest {
24     import std.algorithm : max;
25     ///区間max, 区間加算
26     auto seg = LazySeg!(int, int,
27         (a, b) => max(a, b), (a, b) => a+b, (a, b) => a+b, 0, 0)([2, 1, 4]);
28     
29     //[2, 1, 4]
30     seg[0] = 2; seg[1] = 1; seg[2] = 4;
31     assert(seg[0..3].sum == 4);
32 
33     //[2, 1, 5]
34     seg[2] = 5;
35     assert(seg[0..2].sum == 2);
36     assert(seg[0..3].sum == 5);
37 
38     //[12, 11, 5]
39     seg[0..2] += 10;
40     assert(seg[0..3].sum == 12);
41 }
42 
43 
44 struct LazySegEngine(T, L, alias opTT, alias opTL, alias opLL, T eT, L eL) {
45     import std.typecons : Tuple;
46     alias DataType = T;
47     alias LazyType = L;
48     alias BinSearch = binSearchLazy;
49     alias S = Tuple!(T, "d", L, "lz");
50     uint n, sz, lg;
51     S[] s;
52     this(uint n) {
53         import std.conv : to;
54         import std.algorithm : each;
55         this.n = n;
56         uint lg = 0;
57         while ((2^^lg) < n) lg++;
58         this.lg = lg;
59         sz = 2^^lg;
60         s = new S[](2*sz);
61         s.each!((ref x) => x = S(eT, eL));
62     }
63     this(T[] first) {
64         import std.conv : to;
65         import std.algorithm : each;
66         n = first.length.to!uint;
67         uint lg = 0;
68         while ((2^^lg) < n) lg++;
69         this.lg = lg;
70         sz = 2^^lg;
71 
72         s = new S[](2*sz);
73         s.each!((ref x) => x = S(eT, eL));
74         foreach (i; 0..n) {
75             s[sz+i].d = first[i];
76         }
77         foreach_reverse (i; 1..sz) {
78             update(i);
79         }
80     }
81     @property uint length() const { return n; }
82     pragma(inline):
83     private void lzAdd(uint k, in L x) {
84         s[k].lz = opLL(s[k].lz, x);
85         s[k].d = opTL(s[k].d, x);
86     }
87     public void push(uint k) {
88         if (s[k].lz == eL) return;
89         lzAdd(2*k, s[k].lz);
90         lzAdd(2*k+1, s[k].lz);
91         s[k].lz = eL;
92     }
93     private void update(uint k) {
94         s[k].d = opTT(s[2*k].d, s[2*k+1].d);
95     }
96     T single(uint k) {
97         k += sz;
98         foreach_reverse (uint i; 1..lg+1) {
99             push(k>>i);
100         }
101         return s[k].d;
102     }
103     void singleSet(uint k, T x) {
104         k += sz;
105         foreach_reverse (uint i; 1..lg+1) {
106             push(k>>i);
107         }
108         s[k].d = x;
109         foreach (uint i; 1..lg+1) {
110             update(k>>i);
111         }
112     }
113     T sum(uint a, uint b) {
114         assert(0 <= a && a <= b && b <= n);
115         if (a == b) return eT;
116         a += sz; b--; b += sz;
117         uint tlg = lg;
118         while (true) {
119             uint k = a >> tlg;
120             if (a >> tlg != b >> tlg) {
121                 tlg++;
122                 break;
123             }
124             if (((a-1) >> tlg) + 2 == (b+1) >> tlg) return s[k].d;
125             push(k);
126             tlg--;
127         }
128         T sm = eT;
129         foreach_reverse (l; 0..tlg) {
130             uint k = a >> l;
131             if ((a-1)>>l != a>>l) {
132                 sm = opTT(s[k].d, sm);
133                 break;
134             }
135             push(k);
136             if (!((a >> (l-1)) & 1)) sm = opTT(s[2*k+1].d, sm);
137         }
138         foreach_reverse (l; 0..tlg) {
139             uint k = b >> l;
140             if (b>>l != (b+1)>>l) {
141                 sm = opTT(sm, s[k].d);
142                 break;
143             }
144             push(k);
145             if ((b >> (l-1)) & 1) sm = opTT(sm, s[2*k].d);
146         }
147         return sm;
148     }
149     void add(uint a, uint b, L x) {
150         assert(0 <= a && a <= b && b <= n);
151         if (a == b) return;
152         a += sz; b--; b += sz;
153         uint tlg = lg;
154         while (true) {
155             uint k = a >> tlg;
156             if (a >> tlg != b >> tlg) {
157                 tlg++;
158                 break;
159             }
160             if (((a-1) >> tlg) + 2 == (b+1) >> tlg) {
161                 lzAdd(k, x);
162                 foreach (l; tlg+1..lg+1) {
163                     update(a >> l);
164                 }
165                 return;
166             }
167             push(k);
168             tlg--;
169         }
170         foreach_reverse (l; 0..tlg) {
171             uint k = a >> l;
172             if ((a-1)>>l != a>>l) {
173                 lzAdd(k, x);
174                 foreach (h; l+1..tlg) {
175                     update(a >> h);
176                 }
177                 break;
178             }
179             push(k);
180             if (!((a >> (l-1)) & 1)) lzAdd(2*k+1, x);
181         }
182         foreach_reverse (l; 0..tlg) {
183             uint k = b >> l;
184             if (b>>l != (b+1)>>l) {
185                 lzAdd(k, x);
186                 foreach (h; l+1..tlg) {
187                     update(b >> h);
188                 }
189                 break;
190             }
191             push(k);
192             if ((b >> (l-1)) & 1) lzAdd(2*k, x);
193         }
194         foreach (l; tlg..lg+1) {
195             update(a >> l);
196         }
197     }
198 }
199 
200 
201 
202 unittest {
203     //issue 17466
204     import std.stdio;
205     auto seg = LazySeg!(long[2], long[2],
206         (a, b) => a, (a, b) => a, (a, b) => a, [0L, 0L], [0L, 0L])(10);
207 }
208 
209 
210 int binSearchLazy(bool rev, alias pred, TR)(TR t, int a, int b) {
211     import std.traits : TemplateArgsOf;
212     alias args = TemplateArgsOf!TR;
213     alias opTT = args[2];
214     auto x = args[5];
215     with (t) {
216         static if (!rev) {
217             //left
218             if (pred(x)) return a-1;
219             int pos = a;
220             void f(int a, int b, int l, int r, int k) {
221                 if (b <= l || r <= a) return;
222                 if (a <= l && r <= b && !pred(opTT(x, s[k].d))) {
223                     x = opTT(x, s[k].d);
224                     pos = r;
225                     return;
226                 }
227                 if (l+1 == r) return;
228                 push(k);
229                 int md = (l+r)/2;
230                 f(a, b, l, md, 2*k);
231                 if (pos >= md) f(a, b, md, r, 2*k+1);
232             }
233             f(a, b, 0, sz, 1);
234             return pos;
235         } else {
236             //right
237             if (pred(x)) return b;
238             int pos = b-1;
239             void f(int a, int b, int l, int r, int k) {
240                 if (b <= l || r <= a) return;
241                 if (a <= l && r <= b && !pred(opTT(x, s[k].d))) {
242                     x = opTT(s[k].d, x);
243                     pos = l-1;
244                     return;
245                 }
246                 if (l+1 == r) return;
247                 push(k);
248                 int md = (l+r)/2;
249                 f(a, b, md, r, 2*k+1);
250                 if (pos < md) f(a, b, l, md, 2*k);
251             }
252             f(a, b, 0, sz, 1);
253             return pos;            
254         }
255     }
256 }