manish icon

lazy segment tree

manish | PRO | 03/22/21 07:58:37 PM UTC | 0 ⭐ | 930 👁️ | Never ⏰ | []
Python |

3.63 KB

|

None

|

0 👍

/

0 👎

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

Comments

  •  icon
    01/01/70 12:00:00 AM UTC
    Plain Text |

    0 B

    |

    👍

    /

    👎