func_lazy = lambda a,b,c,d: (a*c, b + a*d) func_data = lambda a,b,val: a * val + b class LazySegmentTree: def __init__(self, data, padding=0, func=lambda a, b: a + b): """initialize the lazy segment tree with data""" self._func = func self._len = len(data) self._size = _size = 1 << (self._len - 1).bit_length() self.data = [padding] * (2 * _size) self.data[_size:_size + self._len] = data for i in reversed(range(1, _size)): self.data[i] = self.data[2 * i] + self.data[2 * i + 1] # TODO self._lazy = [1, 0] * (2 * _size) def _push(self, idx): """push query on idx to its children""" # Let the children know of the queries # TODO idx *= 2 a = self._lazy[idx] b = self._lazy[idx + 1] if (a, b) != (1, 0): b >>= 1 self._lazy[idx] = 1 self._lazy[idx + 1] = 0 self.data[idx] = a * self.data[idx] + b self.data[idx + 1] = a * self.data[idx + 1] + b idx *= 2 self._lazy[idx], self._lazy[idx + 1] = func_lazy(a, b, self._lazy[idx], self._lazy[idx + 1]) self._lazy[idx + 2], self._lazy[idx + 3] = func_lazy(a, b, self._lazy[idx + 2], self._lazy[idx + 3]) def _build(self, idx): """make the changes to idx be known to its ancestors""" idx >>= 1 while idx: # TODO # self._push(idx) a = self._lazy[2 * idx] b = self._lazy[2 * idx + 1] self.data[idx] = a * (self.data[2 * idx] + self.data[2 * idx + 1]) + b idx >>= 1 def _update(self, idx): """updates the node idx to know of all queries applied to it via its ancestors""" for i in reversed(range(1, idx.bit_length())): self._push(idx >> i) def update(self, start, stop, x): self.updateval(start, stop, 0, x) def add(self, start, stop, x): self.updateval(start, stop, 1, x) def updateval(self, start, stop, a, b): """lazily add value to [start, stop)""" stop += 1 start = start_copy = start + self._size stop = stop_copy = stop + self._size # Apply all the lazily stored queries self._update(start) self._update(stop - 1) while start < stop: if start & 1: self.data[start] = a * self.data[start] + b self._lazy[2 * start], self._lazy[2 * start + 1] = func_lazy(a, b, self._lazy[2 * start], self._lazy[2 * start + 1]) start += 1 if stop & 1: stop -= 1 self.data[stop] = a * self.data[stop] + b self._lazy[stop * 2], self._lazy[2 * stop + 1] = func_lazy(a, b, self._lazy[2 * stop], self._lazy[2 * stop + 1]) start >>= 1 stop >>= 1 b <<= 1 self._build(start_copy) self._build(stop_copy - 1) def query(self, start, stop, res=0): stop += 1 start += self._size stop += self._size # Apply all the lazily stored queries self._update(start) self._update(stop - 1) while start < stop: if start & 1: res += self.data[start] start += 1 if stop & 1: stop -= 1 res += self.data[stop] start >>= 1 stop >>= 1 return res def node_size(self, idx): return 1 << self._size.bit_length() - idx.bit_length()
Comments
0 B
|👍
/👎