Segment Tree with Lazy Propagation

Photo by Chris Ried on Unsplash

Segment Tree with Lazy Propagation

·

6 min read

A basic segment tree can efficiently make a range query over the input array and update a single value in the array. If you don't know how to implement a basic segment tree, this article will give you a step-by-step guide on implementing a basic segment tree by using a physical binary tree structure or by using a bounded array.

The basic segment tree works great with point updates (single value updates). However, in many scenarios, we may want to update the values in the given range from left to right, like replacing or incrementing all elements in the range.

One naive solution is that: we can call the modify function for right - left + 1 times. It is obvious that the time complexity of doing so is O(N * log N).

A better idea is that: we don't update a node until it is needed. Then we can avoid unnecessary repeated operations. This technique is called Lazy Propagation. With laze propagation, we can make a range updates in O(log N) time.

The idea of lazy propagation is simple. We will use an auxiliary array called lazy which has the same structure as the segment tree array to mark whether the node has a lazy update. Then we will execute the update when we actually access the node.

Build the tree

Building a segment tree with lazy propagation is the same as building a basic segment tree.

public class LazySegmentTree {
    private final int[] tree;
    private final int[] lazy;
    private final int upperIdx;

    public LazySegmentTree(int[] nums) {
        if (nums == null || nums.length == 0) {
            throw new IllegalArgumentException("input array should not be empty");
        }

        // make enough rooms for the internal nodes.
        this.tree = new int[nums.length * 4];
        this.lazy = new int[nums.length * 4];

        // store the upper index which indicates the end of the range.
        this.upperIdx = nums.length - 1;

        this.buildTree(0, 0, upperIdx, nums);
    }

    private void buildTree(int node, int start, int end, int[] nums) {
        if (start == end) {
            tree[node] = nums[start];
            return;
        }

        int leftChild = 2 * idx + 1;
        int rightChild = 2 * idx + 2;
        int mid = start + (end - start) / 2;

        buildTree(leftChild, start, mid, nums);
        buildTree(rightChild, mid + 1, end, nums);

        tree[node] = tree[leftChild] + tree[rightChild];
    }
}

As we can see, building a lazy segment tree is the same as building a basic one. So the time complexity is O(N).

Modify values in a range

When updating a single value in a segment tree, we update the corresponding leaf node and then recursively update all its parent nodes. However, when updating a range of values, we need to implement a slightly different approach.

Suppose we need to update node X in the segment tree. Instead of updating all the nodes in the range, we only update node X and mark its children for future updates. To achieve this, we maintain another binary tree called the lazy tree, which has the same structure as the segment tree.

We initialize the lazy tree with all zeros, indicating that there are no pending updates. When we perform a range update, we mark the corresponding nodes in the lazy tree with non-zero values, indicating that there are pending updates for those nodes.

Later, when we access a node with a non-zero lazy value, we apply the pending updates to that node and then notify its children of the update. This way, we can efficiently update a range of nodes in the segment tree without having to update all the nodes in the range at once.

public class LazySegmentTree {
    // ... code for building the tree

    public void modifyRange(int left, int right, int val) {
        if (left > right) {
            throw new IllegalArgumentException("left should not be greater than right");
        }
        modify(0, 0, upperIdx, left, right, val);
    }

    private void modify(int node, int start, int end, int left, int right, int val) {
        lazyUpdate(node, start, end);

        // If the node lies outside the input range,
        // then we simply return.
        if (right < start || end < left) {
            return;
        }

        // If the node lies fully inside the input range,
        // then we simply update the node and mark its children lazy.
        if (left <= start && end <= right) {
            tree[node] += (end - start + 1) * val;
            if (start < end) {
                lazy[2 * node + 1] += val;
                lazy[2 * node + 2] += val;
            }
            return;
        }

        int leftChild = 2 * idx + 1;
        int rightChild = 2 * idx + 2;
        int mid = start + (end - start) / 2;

        // recursively update the left child.
        update(leftChild, start, mid, left, right, val);
        // recursively update the right child.
        update(rightChild, mid + 1, end, left, right, val);

        // Updating the parent with the values of the left and right child.
        tree[node] = tree[leftChild] + tree[rightChild];
    }

    private void lazyUpdate(int node, int start, int end) {
        // zero means there are no pending updates.
        if (lazy[node] == 0) {
            return;
        }

        // apply the lazy range updates to the node.
        tree[node] += (end - start + 1) * lazy[node];

        // if its children exist then mark them as lazy.
        if (start < end) {
            lazy[2 * node + 1] += lazy[node];
            lazy[2 * node + 2] += lazy[node];
        }

        // mark the node no longer lazy.
        lazy[node] = 0;
    }
}

As we can see, with lazy propagation, when updating values in a given range, the longest path we may traverse is from the root node to the leaf node. So the time complexity of modifying a single value is O(height) ~= O(log N).

Range query function

When making a range query in the segment tree, it is quite similar to the range update: we check if the current node has a non-zero lazy value. If so, we need to update the node before proceeding with the query. We also mark its child nodes as lazy if they exist. Then we return the node value.

public class LazySegmentTree {
    // ... code for building the tree & range updating

    public int queryRange(int left, int right) {
        if (left > right) {
            throw new IllegalArgumentException("left should not be greater than right");
        }
        return query(0, 0, upperIdx, left, right);
    }

    private int query(int node, int start, int end, int left, int right) {
        lazyUpdate(node, start, end);

        // If the node lies outside the input range,
        // then we simply return 0.
        if (right < start || end < left) {
            return 0;
        }

        // If the node lies fully inside the input range,
        // then we simply return the node value.
        if (left <= start && end <= right) {
            return tree[node];
        }

        int mid = start + (end - start) / 2;

        // query the left child.
        int leftRes = query(2 * node + 1, start, mid, left, right);
        // query the right child.
        int rightRes = query(2 * node + 2, mid + 1, end, left, right);

        return leftRes + rightRes;
    }
}

As we can see, making a range query in a lazy segment tee is almost the same as making a range query in a basic segment tree. The biggest difference is that: we do a lazy update first, which takes O(1) time. So the time complexity of this operation is O(log N).

Conclusion

With the lazy propagation mechanism, the segment tree has evolved to support both range queries and range updates in O(log N) time. This enables the segment tree to answer queries over an array and update the array in a minimum time.

In the next article, we will see how to use the lazy segment tree to solve a practical problem.