Advertisement
manish

Untitled

Mar 22nd, 2021
591
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 3.32 KB | None | 0 0
  1. class LazySegmentTree:
  2.     def __init__(self, data, padding = 0, func=lambda a,b: a + b):
  3.         """initialize the lazy segment tree with data"""
  4.         self._func = func
  5.         self._len = len(data)
  6.         self._size = _size = 1 << (self._len - 1).bit_length()
  7.        
  8.         self.data = [padding] * (2 * _size)
  9.         self.data[_size:_size + self._len] = data
  10.         for i in reversed(range(1, _size)):
  11.             self.data[i] = self.data[2 * i] + self.data[2 * i + 1]    
  12.         #TODO
  13.         self._lazy = [1,0] * (2 * _size)
  14.  
  15.     def _push(self, idx):
  16.         """push query on idx to its children"""
  17.         # Let the children know of the queries
  18.         # TODO
  19.         idx *= 2
  20.         a = self._lazy[idx]
  21.         b = self._lazy[idx + 1]
  22.         if (a,b) != (1,0):
  23.             b >>= 1
  24.             self._lazy[idx] = 1
  25.             self._lazy[idx+1] = 0
  26.            
  27.             self.data[idx] = a * self.data[idx] + b
  28.             self.data[idx + 1] = a * self.data[ idx + 1] + b
  29.            
  30.             idx *= 2
  31.             self._lazy[idx], self._lazy[idx + 1] = func_lazy((a,b), (self._lazy[idx], self._lazy[idx+1]))
  32.             self._lazy[idx + 2], self._lazy[idx + 3] = func_lazy((a,b), (self._lazy[idx + 2], self._lazy[idx + 3]))
  33.     def _build(self, idx):
  34.         """make the changes to idx be known to its ancestors"""
  35.         idx >>= 1
  36.         while idx:
  37.             # TODO
  38.             #self._push(idx)
  39.             a = self._lazy[2 * idx]
  40.             b = self._lazy[2 * idx + 1]
  41.             self.data[idx] = a * (self.data[2 * idx] + self.data[2 * idx + 1]) + b
  42.             idx >>= 1
  43.  
  44.     def _update(self, idx):
  45.         """updates the node idx to know of all queries applied to it via its ancestors"""
  46.         for i in reversed(range(1, idx.bit_length())):
  47.             self._push(idx >> i)
  48.  
  49.     def add(self, start, stop, (a,b)):
  50.         """lazily add value to [start, stop)"""
  51.         start = start_copy = start + self._size
  52.         stop = stop_copy = stop + self._size
  53.  
  54.         # Apply all the lazily stored queries
  55.         self._update(start); self._update(stop - 1)
  56.  
  57.         while start < stop:
  58.             if start & 1:
  59.                 self.data[start] = a * self.data[start] + b
  60.                 self._lazy[2*start],self._lazy[2*start+1] = func_lazy((a,b), (self._lazy[2 * start], self._lazy[2 * start + 1]))
  61.                 start += 1
  62.             if stop & 1:
  63.                 stop -= 1
  64.                 self.data[stop] = a * self.data[stop] + b
  65.                 self._lazy[stop * 2], self._lazy[2 * stop + 1] = func_lazy((a,b), (self._lazy[2 * stop], self._lazy[2 * stop + 1]))
  66.             start >>= 1; stop >>= 1; b <<= 1
  67.        
  68.         self._build(start_copy); self._build(stop_copy - 1)
  69.  
  70.     def query(self, start, stop, res = 0):
  71.         """func of data[start, stop)"""
  72.         start += self._size; stop += self._size
  73.  
  74.         # Apply all the lazily stored queries
  75.         self._update(start); self._update(stop - 1)
  76.         while start < stop:
  77.             if start & 1:
  78.                 res += self.data[start]
  79.                 start += 1
  80.             if stop & 1:
  81.                 stop -= 1
  82.                 res += self.data[stop]
  83.             start >>= 1; stop >>= 1
  84.         return res
  85.  
  86.     def node_size(self, idx):
  87.         return 1 << self._size.bit_length() - idx.bit_length()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement