Pasted

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()