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 add(self, start, stop, (a,b)): """lazily add value to [start, stop)""" 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): """func of data[start, stop)""" 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()