diff --git a/lazysegtree.py b/lazysegtree.py index 608930a..c979754 100644 --- a/lazysegtree.py +++ b/lazysegtree.py @@ -23,111 +23,254 @@ def __init__(self, V, OP, E, MAPPING, COMPOSITION, ID): self.mapping = MAPPING self.composition = COMPOSITION self.identity = ID + d = self.d for i in range(self.n): - self.d[self.size + i] = V[i] + d[self.size + i] = V[i] + op = self.op for i in range(self.size - 1, 0, -1): - self.update(i) + d[i] = op(d[2 * i], d[2 * i + 1]) def set(self, p, x): assert 0 <= p and p < self.n - p += self.size + d = self.d + lz = self.lz + mapping = self.mapping + composition = self.composition + identity = self.identity + op = self.op + size = self.size + p += size for i in range(self.log, 0, -1): - self.push(p >> i) - self.d[p] = x + k = p >> i + k2 = 2 * k + pf = lz[k] + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + k2 += 1 + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + lz[k] = identity + d[p] = x for i in range(1, self.log + 1): - self.update(p >> i) + k = p >> i + d[k] = op(d[2 * k], d[2 * k + 1]) def get(self, p): assert 0 <= p and p < self.n - p += self.size + d = self.d + lz = self.lz + mapping = self.mapping + composition = self.composition + identity = self.identity + size = self.size + p += size for i in range(self.log, 0, -1): - self.push(p >> i) - return self.d[p] + k = p >> i + k2 = 2 * k + pf = lz[k] + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + k2 += 1 + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + lz[k] = identity + return d[p] def prod(self, l, r): assert 0 <= l and l <= r and r <= self.n if l == r: return self.e - l += self.size - r += self.size - for i in range(self.log, 0, -1): + d = self.d + lz = self.lz + mapping = self.mapping + composition = self.composition + identity = self.identity + op = self.op + size = self.size + log = self.log + l += size + r += size + for i in range(log, 0, -1): if ((l >> i) << i) != l: - self.push(l >> i) + k = l >> i + k2 = 2 * k + pf = lz[k] + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + k2 += 1 + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + lz[k] = identity if ((r >> i) << i) != r: - self.push(r >> i) + k = r >> i + k2 = 2 * k + pf = lz[k] + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + k2 += 1 + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + lz[k] = identity sml, smr = self.e, self.e while l < r: if l & 1: - sml = self.op(sml, self.d[l]) + sml = op(sml, d[l]) l += 1 if r & 1: r -= 1 - smr = self.op(self.d[r], smr) + smr = op(d[r], smr) l >>= 1 r >>= 1 - return self.op(sml, smr) + return op(sml, smr) def all_prod(self): return self.d[1] def apply_point(self, p, f): assert 0 <= p and p < self.n - p += self.size + d = self.d + lz = self.lz + mapping = self.mapping + composition = self.composition + identity = self.identity + op = self.op + size = self.size + p += size for i in range(self.log, 0, -1): - self.push(p >> i) - self.d[p] = self.mapping(f, self.d[p]) + k = p >> i + k2 = 2 * k + pf = lz[k] + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + k2 += 1 + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + lz[k] = identity + d[p] = mapping(f, d[p]) for i in range(1, self.log + 1): - self.update(p >> i) + k = p >> i + d[k] = op(d[2 * k], d[2 * k + 1]) def apply(self, l, r, f): assert 0 <= l and l <= r and r <= self.n if l == r: return - l += self.size - r += self.size - for i in range(self.log, 0, -1): + d = self.d + lz = self.lz + mapping = self.mapping + composition = self.composition + identity = self.identity + op = self.op + size = self.size + log = self.log + l += size + r += size + for i in range(log, 0, -1): if ((l >> i) << i) != l: - self.push(l >> i) + k = l >> i + k2 = 2 * k + pf = lz[k] + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + k2 += 1 + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + lz[k] = identity if ((r >> i) << i) != r: - self.push((r - 1) >> i) + k = (r - 1) >> i + k2 = 2 * k + pf = lz[k] + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + k2 += 1 + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + lz[k] = identity l2, r2 = l, r while l < r: if l & 1: - self.all_apply(l, f) + d[l] = mapping(f, d[l]) + if l < size: + lz[l] = composition(f, lz[l]) l += 1 if r & 1: r -= 1 - self.all_apply(r, f) + d[r] = mapping(f, d[r]) + if r < size: + lz[r] = composition(f, lz[r]) l >>= 1 r >>= 1 l, r = l2, r2 - for i in range(1, self.log + 1): + for i in range(1, log + 1): if ((l >> i) << i) != l: - self.update(l >> i) + k = l >> i + d[k] = op(d[2 * k], d[2 * k + 1]) if ((r >> i) << i) != r: - self.update((r - 1) >> i) + k = (r - 1) >> i + d[k] = op(d[2 * k], d[2 * k + 1]) def max_right(self, l, g): assert 0 <= l and l <= self.n assert g(self.e) if l == self.n: return self.n - l += self.size + d = self.d + lz = self.lz + mapping = self.mapping + composition = self.composition + identity = self.identity + op = self.op + size = self.size + l += size for i in range(self.log, 0, -1): - self.push(l >> i) + k = l >> i + k2 = 2 * k + pf = lz[k] + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + k2 += 1 + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + lz[k] = identity sm = self.e while 1: while l % 2 == 0: l >>= 1 - if not (g(self.op(sm, self.d[l]))): - while l < self.size: - self.push(l) + if not (g(op(sm, d[l]))): + while l < size: + k = l + k2 = 2 * k + pf = lz[k] + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + k2 += 1 + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + lz[k] = identity l = 2 * l - if g(self.op(sm, self.d[l])): - sm = self.op(sm, self.d[l]) + if g(op(sm, d[l])): + sm = op(sm, d[l]) l += 1 - return l - self.size - sm = self.op(sm, self.d[l]) + return l - size + sm = op(sm, d[l]) l += 1 if (l & -l) == l: break @@ -138,23 +281,50 @@ def min_left(self, r, g): assert g(self.e) if r == 0: return 0 - r += self.size + d = self.d + lz = self.lz + mapping = self.mapping + composition = self.composition + identity = self.identity + op = self.op + size = self.size + r += size for i in range(self.log, 0, -1): - self.push((r - 1) >> i) + k = (r - 1) >> i + k2 = 2 * k + pf = lz[k] + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + k2 += 1 + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + lz[k] = identity sm = self.e while 1: r -= 1 while r > 1 and (r % 2): r >>= 1 - if not (g(self.op(self.d[r], sm))): - while r < self.size: - self.push(r) + if not (g(op(d[r], sm))): + while r < size: + k = r + k2 = 2 * k + pf = lz[k] + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + k2 += 1 + d[k2] = mapping(pf, d[k2]) + if k2 < size: + lz[k2] = composition(pf, lz[k2]) + lz[k] = identity r = 2 * r + 1 - if g(self.op(self.d[r], sm)): - sm = self.op(self.d[r], sm) + if g(op(d[r], sm)): + sm = op(d[r], sm) r -= 1 - return r + 1 - self.size - sm = self.op(self.d[r], sm) + return r + 1 - size + sm = op(d[r], sm) if (r & -r) == r: break return 0 diff --git a/segtree.py b/segtree.py index c32257e..5738376 100644 --- a/segtree.py +++ b/segtree.py @@ -13,17 +13,22 @@ def __init__(self, V, OP, E): self.log = (self.n - 1).bit_length() self.size = 1 << self.log self.d = [E for i in range(2 * self.size)] + d = self.d for i in range(self.n): - self.d[self.size + i] = V[i] + d[self.size + i] = V[i] + op = self.op for i in range(self.size - 1, 0, -1): - self.update(i) + d[i] = op(d[2 * i], d[2 * i + 1]) def set(self, p, x): assert 0 <= p and p < self.n + d = self.d + op = self.op p += self.size - self.d[p] = x + d[p] = x for i in range(1, self.log + 1): - self.update(p >> i) + k = p >> i + d[k] = op(d[2 * k], d[2 * k + 1]) def get(self, p): assert 0 <= p and p < self.n @@ -31,20 +36,22 @@ def get(self, p): def prod(self, l, r): assert 0 <= l and l <= r and r <= self.n + d = self.d + op = self.op sml = self.e smr = self.e l += self.size r += self.size while l < r: if l & 1: - sml = self.op(sml, self.d[l]) + sml = op(sml, d[l]) l += 1 if r & 1: - smr = self.op(self.d[r - 1], smr) + smr = op(d[r - 1], smr) r -= 1 l >>= 1 r >>= 1 - return self.op(sml, smr) + return op(sml, smr) def all_prod(self): return self.d[1] @@ -54,19 +61,22 @@ def max_right(self, l, f): assert f(self.e) if l == self.n: return self.n - l += self.size + d = self.d + op = self.op + size = self.size + l += size sm = self.e while 1: while l % 2 == 0: l >>= 1 - if not (f(self.op(sm, self.d[l]))): - while l < self.size: + if not (f(op(sm, d[l]))): + while l < size: l = 2 * l - if f(self.op(sm, self.d[l])): - sm = self.op(sm, self.d[l]) + if f(op(sm, d[l])): + sm = op(sm, d[l]) l += 1 - return l - self.size - sm = self.op(sm, self.d[l]) + return l - size + sm = op(sm, d[l]) l += 1 if (l & -l) == l: break @@ -77,20 +87,23 @@ def min_left(self, r, f): assert f(self.e) if r == 0: return 0 - r += self.size + d = self.d + op = self.op + size = self.size + r += size sm = self.e while 1: r -= 1 while r > 1 and (r % 2): r >>= 1 - if not (f(self.op(self.d[r], sm))): - while r < self.size: + if not (f(op(d[r], sm))): + while r < size: r = 2 * r + 1 - if f(self.op(self.d[r], sm)): - sm = self.op(self.d[r], sm) + if f(op(d[r], sm)): + sm = op(d[r], sm) r -= 1 - return r + 1 - self.size - sm = self.op(self.d[r], sm) + return r + 1 - size + sm = op(d[r], sm) if (r & -r) == r: break return 0