Python

Dynamic Segment Tree

Python 2 min read

Why Dynamic Segment Tree?

As know for segment tree, the space complexity is is up to 4n, which n is upper bound of data range. However, :-(, n some times becomes very large, such as $10^9$ or infinity. OOM problem would occur if we allocate 4n size array for query and update options.

The Dynamic Segment Tree (DST) collects many excellent features. The first one is lazy initiation, part of space is allocated while it's being needed. The other features are similar to segment tree but superior to it.

Dynamic Segment Tree implementation

The implementation of python are shown as follows:

class Node:
    def __init__(self, start, end):
        self.start = start
        self.end = end
        self.total = 0
        self.lazy = 0
        self.left = None
        self.right = None

class DynamicSegmentTree:
    def __init__(self, start, end):
        self.root = Node(start, end)

    def updateRange(self, node, start, end, val):
        if node.lazy != 0:
            node.total += (node.end - node.start + 1) * node.lazy
            if node.start != node.end:
                if not node.left:
                    node.left = Node(node.start, (node.start + node.end) // 2)
                if not node.right:
                    node.right = Node((node.start + node.end) // 2 + 1, node.end)
                node.left.lazy += node.lazy
                node.right.lazy += node.lazy
            node.lazy = 0

        if node.start > end or node.end < start:
            return

        if node.start >= start and node.end <= end:
            node.total += (node.end - node.start + 1) * val
            if node.start != node.end:
                if not node.left:
                    node.left = Node(node.start, (node.start + node.end) // 2)
                if not node.right:
                    node.right = Node((node.start + node.end) // 2 + 1, node.end)
                node.left.lazy += val
                node.right.lazy += val
            return

        mid = (node.start + node.end) // 2
        if not node.left:
            node.left = Node(node.start, mid)
        if not node.right:
            node.right = Node(mid + 1, node.end)
        self.updateRange(node.left, start, end, val)
        self.updateRange(node.right, start, end, val)

        node.total = node.left.total + node.right.total

    def queryRange(self, node, start, end):
        if not node or start > node.end or end < node.start:
            return 0

        if node.lazy != 0:
            node.total += (node.end - node.start + 1) * node.lazy
            if node.start != node.end:
                if not node.left:
                    node.left = Node(node.start, (node.start + node.end) // 2)
                if not node.right:
                    node.right = Node((node.start + node.end) // 2 + 1, node.end)
                node.left.lazy += node.lazy
                node.right.lazy += node.lazy
            node.lazy = 0

        if start <= node.start and end >= node.end:
            return node.total

        return self.queryRange(node.left, start, end) + self.queryRange(node.right, start, end)

# Usage
tree = DynamicSegmentTree(0, 1000)  # Range 0 to 1000
tree.updateRange(tree.root, 10, 20, 5)  # Update range 10 to 20 by adding 5
print(tree.queryRange(tree.root, 0, 15))  # Query sum from index 0 to 50

Read more