Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- 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()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement