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