Segment Tree for Range Assignment and Range Sum Queries

Given an array of size N filled with all 0s, the task is to answer Q queries, where the queries can be one of the two types:

• Type 1 (1, L, R, X): Assign value X to all elements on the segment from L to R−1, and
• Type 2 (2, L, R): Find the sum on the segment from L to R−1.

Examples:

Input: N = 5, Q = 3, queries[][] = {{1, 0, 3, 5}, {2, 1, 4}, {1, 2, 4, 10}}
Output: 10
Explanation:

• Initially the array is {0, 0, 0, 0, 0}
• First query is to assign value 5 from index 0 to 2, so after first query the array is {5, 5, 5, 0, 0}.
• Second query is to find the sum of segment from index 1 to 3, so the sum = 5 + 5 = 10.
• Third query is to assign value 10 from index 2 to 3, so after third query the array is {5, 5, 10, 10, 0}.

Input: N = 10, Q = 4, queries[q][4] = {{1, 0, 5, 3}, {2, 2, 7}, {1, 4, 9, 2}, {2, 0, 9}}
Output: 9 22
Explanation:

• Initially the array is {0, 0, 0, 0, 0, 0, 0, 0, 0, 0}
• First query is to assign value 3 from index 0 to 4, so after first query the array is {3, 3, 3, 3, 3, 0, 0, 0, 0, 0}.
• Second query is to find the sum of segment from index 2 to 6, so the sum = 3 + 3 + 3 + 0 + 0 = 9.
• Third query is to assign value 2 from index 4 to 8, so after third query the array is {3, 3, 3, 3, 2, 2, 2, 2, 2, 0}.
• Fourth query is to find the sum of segment from index 0 to 8, so the sum = 3 + 3 + 3 + 3 + 2 + 2 + 2 + 2 = 22.

Approach: To solve the problem, follow the below idea:

The main idea is to build a segment tree and use lazy propagation to efficiently handle the range update and range sum queries. The segment tree is built in such a way that all levels of the tree are fully filled except possibly for the last level, which is filled from left to right. This makes the height of the tree log(n), leading to efficient query and update operations.

Step-by-step algorithm:

• Declare arrays tree and lazy for the segment tree and lazy propagation.
• Implement the updateRange function to update a range in the segment tree.
  • Handle pending updates using the lazy array.
  • If the current segment is fully in range, update the node and propagate the information to its children.
  • If not completely in range but overlaps, recursively update the left and right children and use their results to update the current node.
• Implement the queryRange function to calculate the sum on a given range.
  • Handle pending updates using the lazy array.
  • If the current segment is fully in range, return the value of the segment.
  • If not completely in range but overlaps, recursively query the left and right children and return the sum of their results.
• Return the final answer of every query of type 2.

Below is the implementation of the algorithm:

C++

#include <bits/stdc++.h>
using namespace std;

const int MAX = 1e5; // Max size of array
int tree[4 * MAX] = { 0 }; // Segment tree
int lazy[4 * MAX] = { 0 }; // Lazy array

// Function to update the segment tree
void updateRange(int node, int start, int end, int l, int r,
                int val)
{

    // If lazy[node] is non-zero, then there are some
    // pending updates. So we need to make sure the node is
    // updated
    if (lazy[node] != 0) {

        // Updating the node
        tree[node] = (end - start + 1) * lazy[node];

        // Passing the update information to its children
        if (start != end) {
            lazy[node * 2] = lazy[node];
            lazy[node * 2 + 1] = lazy[node];
        }

        // Resetting the lazy value for the current node
        lazy[node] = 0;
    }

    // Out of range
    if (start > end or start > r or end < l)
        return;

    // Current segment is fully in range
    if (start >= l and end <= r) {

        // Update the node
        tree[node] = (end - start + 1) * val;

        // Pass the update information to its children
        if (start != end) {
            lazy[node * 2] = val;
            lazy[node * 2 + 1] = val;
        }
        return;
    }

    // If not completely in range but overlaps, recur for
    // children,

    int mid = (start + end) / 2;
    updateRange(node * 2, start, mid, l, r, val);
    updateRange(node * 2 + 1, mid + 1, end, l, r, val);

    // Use the result of children calls to update this node
    tree[node] = tree[node * 2] + tree[node * 2 + 1];
}

// Function to calculate the sum on a given range
int queryRange(int node, int start, int end, int l, int r)
{
    // Out of range
    if (start > end or start > r or end < l)
        return 0;

    // If there are pending updates
    if (lazy[node] != 0) {

        // Updating the node
        tree[node] = (end - start + 1) * lazy[node];

        // Passing the update information to its children
        if (start != end) {
            lazy[node * 2] = lazy[node];
            lazy[node * 2 + 1] = lazy[node];
        }

        // Resetting the lazy value for the current node
        lazy[node] = 0;
    }

    // At this point we are sure that pending lazy updates
    // are done for current node. So we can return value
    if (start >= l and end <= r)
        return tree[node];

    // If not completely in range but overlaps, recur for
    // children,
    int mid = (start + end) / 2;
    int p1 = queryRange(node * 2, start, mid, l, r);
    int p2 = queryRange(node * 2 + 1, mid + 1, end, l, r);

    // Use the result of children calls to update this node
    return (p1 + p2);
}

int main()
{
    // Hardcoded input
    int n = 5, q = 3;
    int queries[q][4]
        = { { 1, 0, 3, 5 }, { 2, 1, 4 }, { 1, 2, 4, 10 } };

    for (int i = 0; i < q; i++) {
        int type = queries[i][0];

        if (type == 1) {
            int l = queries[i][1], r = queries[i][2],
                v = queries[i][3];
            updateRange(1, 0, n - 1, l, r - 1, v);
        }

        else {
            int l = queries[i][1], r = queries[i][2];
            cout << queryRange(1, 0, n - 1, l, r - 1)
                << "\n";
        }
    }
    return 0;
}

Java

import java.util.Arrays;

class Main {

    static final int MAX = 100000;
    static int[] tree = new int[4 * MAX];
    static int[] lazy = new int[4 * MAX];

    // Function to update the segment tree
    static void updateRange(int node, int start, int end, int l, int r, int val) {

        // If lazy[node] is non-zero, then there are some
        // pending updates. So we need to make sure the node is
        // updated
        if (lazy[node] != 0) {

            // Updating the node
            tree[node] = (end - start + 1) * lazy[node];

            // Passing the update information to its children
            if (start != end) {
                lazy[node * 2] = lazy[node];
                lazy[node * 2 + 1] = lazy[node];
            }

            // Resetting the lazy value for the current node
            lazy[node] = 0;
        }

        // Out of range
        if (start > end || start > r || end < l)
            return;

        // Current segment is fully in range
        if (start >= l && end <= r) {

            // Update the node
            tree[node] = (end - start + 1) * val;

            // Pass the update information to its children
            if (start != end) {
                lazy[node * 2] = val;
                lazy[node * 2 + 1] = val;
            }
            return;
        }

        // If not completely in range but overlaps, recur for children,
        int mid = (start + end) / 2;
        updateRange(node * 2, start, mid, l, r, val);
        updateRange(node * 2 + 1, mid + 1, end, l, r, val);

        // Use the result of children calls to update this node
        tree[node] = tree[node * 2] + tree[node * 2 + 1];
    }

    // Function to calculate the sum on a given range
    static int queryRange(int node, int start, int end, int l, int r) {

        // Out of range
        if (start > end || start > r || end < l)
            return 0;

        // If there are pending updates
        if (lazy[node] != 0) {

            // Updating the node
            tree[node] = (end - start + 1) * lazy[node];

            // Passing the update information to its children
            if (start != end) {
                lazy[node * 2] = lazy[node];
                lazy[node * 2 + 1] = lazy[node];
            }

            // Resetting the lazy value for the current node
            lazy[node] = 0;
        }

        // At this point, we are sure that pending lazy updates
        // are done for the current node. So we can return the value
        if (start >= l && end <= r)
            return tree[node];

        // If not completely in range but overlaps, recur for children,
        int mid = (start + end) / 2;
        int p1 = queryRange(node * 2, start, mid, l, r);
        int p2 = queryRange(node * 2 + 1, mid + 1, end, l, r);

        // Use the result of children calls to update this node
        return (p1 + p2);
    }

    public static void main(String[] args) {

        // Hardcoded input
        int n = 5, q = 3;
        int[][] queries = { { 1, 0, 3, 5 }, { 2, 1, 4 }, { 1, 2, 4, 10 } };

        for (int i = 0; i < q; i++) {
            int type = queries[i][0];

            if (type == 1) {
                int l = queries[i][1], r = queries[i][2], v = queries[i][3];
                updateRange(1, 0, n - 1, l, r - 1, v);
            } else {
                int l = queries[i][1], r = queries[i][2];
                System.out.println(queryRange(1, 0, n - 1, l, r - 1));
            }
        }
    }
}

// This code is contributed by shivamgupta310570

Python3

# Python code

MAX = 10**5  # Max size of array
tree = [0] * (4 * MAX)  # Segment tree
lazy = [0] * (4 * MAX)  # Lazy array

# Function to update the segment tree
def updateRange(node, start, end, l, r, val):
    # If lazy[node] is non-zero, then there are some
    # pending updates. So we need to make sure the node is
    # updated
    if lazy[node] != 0:
        # Updating the node
        tree[node] = (end - start + 1) * lazy[node]

        # Passing the update information to its children
        if start != end:
            lazy[node * 2] = lazy[node]
            lazy[node * 2 + 1] = lazy[node]

        # Resetting the lazy value for the current node
        lazy[node] = 0

    # Out of range
    if start > end or start > r or end < l:
        return

    # Current segment is fully in range
    if start >= l and end <= r:
        # Update the node
        tree[node] = (end - start + 1) * val

        # Pass the update information to its children
        if start != end:
            lazy[node * 2] = val
            lazy[node * 2 + 1] = val
        return

    # If not completely in range but overlaps, recur for children
    mid = (start + end) // 2
    updateRange(node * 2, start, mid, l, r, val)
    updateRange(node * 2 + 1, mid + 1, end, l, r, val)

    # Use the result of children calls to update this node
    tree[node] = tree[node * 2] + tree[node * 2 + 1]

# Function to calculate the sum on a given range
def queryRange(node, start, end, l, r):
    # Out of range
    if start > end or start > r or end < l:
        return 0

    # If there are pending updates
    if lazy[node] != 0:
        # Updating the node
        tree[node] = (end - start + 1) * lazy[node]

        # Passing the update information to its children
        if start != end:
            lazy[node * 2] = lazy[node]
            lazy[node * 2 + 1] = lazy[node]

        # Resetting the lazy value for the current node
        lazy[node] = 0

    # At this point, we are sure that pending lazy updates
    # are done for the current node. So we can return the value
    if start >= l and end <= r:
        return tree[node]

    # If not completely in range but overlaps, recur for children
    mid = (start + end) // 2
    p1 = queryRange(node * 2, start, mid, l, r)
    p2 = queryRange(node * 2 + 1, mid + 1, end, l, r)

    # Use the result of children calls to update this node
    return p1 + p2

# Hardcoded input
n = 5
q = 3
queries = [[1, 0, 3, 5], [2, 1, 4], [1, 2, 4, 10]]

for i in range(q):
    type = queries[i][0]

    if type == 1:
        l, r, v = queries[i][1], queries[i][2], queries[i][3]
        updateRange(1, 0, n - 1, l, r - 1, v)

    else:
        l, r = queries[i][1], queries[i][2]
        print(queryRange(1, 0, n - 1, l, r - 1))

# This code is contributed by shivamgupta310570

C#

using System;

public class SegmentTree {
    const int MAX = 100000; // Max size of array
    int[] tree = new int[4 * MAX]; // Segment tree
    int[] lazy = new int[4 * MAX]; // Lazy array

    // Function to update the segment tree
    public void UpdateRange(int node, int start, int end,
                            int l, int r, int val)
    {
        // If lazy[node] is non-zero, then there are some
        // pending updates. So we need to make sure the node
        // is updated
        if (lazy[node] != 0) {
            // Updating the node
            tree[node] = (end - start + 1) * lazy[node];

            // Passing the update information to its
            // children
            if (start != end) {
                lazy[node * 2] = lazy[node];
                lazy[node * 2 + 1] = lazy[node];
            }

            // Resetting the lazy value for the current node
            lazy[node] = 0;
        }

        // Out of range
        if (start > end || start > r || end < l)
            return;

        // Current segment is fully in range
        if (start >= l && end <= r) {
            // Update the node
            tree[node] = (end - start + 1) * val;

            // Pass the update information to its children
            if (start != end) {
                lazy[node * 2] = val;
                lazy[node * 2 + 1] = val;
            }
            return;
        }

        // If not completely in range but overlaps, recur
        // for children,

        int mid = (start + end) / 2;
        UpdateRange(node * 2, start, mid, l, r, val);
        UpdateRange(node * 2 + 1, mid + 1, end, l, r, val);

        // Use the result of children calls to update this
        // node
        tree[node] = tree[node * 2] + tree[node * 2 + 1];
    }

    // Function to calculate the sum on a given range
    public int QueryRange(int node, int start, int end,
                          int l, int r)
    {
        // Out of range
        if (start > end || start > r || end < l)
            return 0;

        // If there are pending updates
        if (lazy[node] != 0) {
            // Updating the node
            tree[node] = (end - start + 1) * lazy[node];

            // Passing the update information to its
            // children
            if (start != end) {
                lazy[node * 2] = lazy[node];
                lazy[node * 2 + 1] = lazy[node];
            }

            // Resetting the lazy value for the current node
            lazy[node] = 0;
        }

        // At this point we are sure that pending lazy
        // updates are done for current node. So we can
        // return value
        if (start >= l && end <= r)
            return tree[node];

        // If not completely in range but overlaps, recur
        // for children,
        int mid = (start + end) / 2;
        int p1 = QueryRange(node * 2, start, mid, l, r);
        int p2
            = QueryRange(node * 2 + 1, mid + 1, end, l, r);

        // Use the result of children calls to update this
        // node
        return (p1 + p2);
    }

    public static void Main(string[] args)
    {
        SegmentTree segmentTree = new SegmentTree();

        // Hardcoded input
        int n = 5, q = 3;
        int[][] queries = new int[q][];
        queries[0] = new int[] { 1, 0, 3, 5 };
        queries[1] = new int[] { 2, 1, 4 };
        queries[2] = new int[] { 1, 2, 4, 10 };

        for (int i = 0; i < q; i++) {
            int type = queries[i][0];

            if (type == 1) {
                int l = queries[i][1], r = queries[i][2],
                    v = queries[i][3];
                segmentTree.UpdateRange(1, 0, n - 1, l,
                                        r - 1, v);
            }
            else {
                int l = queries[i][1], r = queries[i][2];
                Console.WriteLine(segmentTree.QueryRange(
                    1, 0, n - 1, l, r - 1));
            }
        }
    }
}

Javascript

// JavaScript Implementation

const MAX = 1e5; // Max size of array
let tree = Array(4 * MAX).fill(0); // Segment tree
let lazy = Array(4 * MAX).fill(0); // Lazy array

// Function to update the segment tree
function updateRange(node, start, end, l, r, val) {

    // If lazy[node] is non-zero, then there are some
    // pending updates. So we need to make sure the node is
    // updated
    if (lazy[node] !== 0) {

        // Updating the node
        tree[node] = (end - start + 1) * lazy[node];

        // Passing the update information to its children
        if (start !== end) {
            lazy[node * 2] = lazy[node];
            lazy[node * 2 + 1] = lazy[node];
        }

        // Resetting the lazy value for the current node
        lazy[node] = 0;
    }

    // Out of range
    if (start > end || start > r || end < l)
        return;

    // Current segment is fully in range
    if (start >= l && end <= r) {

        // Update the node
        tree[node] = (end - start + 1) * val;

        // Pass the update information to its children
        if (start !== end) {
            lazy[node * 2] = val;
            lazy[node * 2 + 1] = val;
        }
        return;
    }

    // If not completely in range but overlaps, recur for
    // children,

    let mid = Math.floor((start + end) / 2);
    updateRange(node * 2, start, mid, l, r, val);
    updateRange(node * 2 + 1, mid + 1, end, l, r, val);

    // Use the result of children calls to update this node
    tree[node] = tree[node * 2] + tree[node * 2 + 1];
}

// Function to calculate the sum on a given range
function queryRange(node, start, end, l, r) {
    // Out of range
    if (start > end || start > r || end < l)
        return 0;

    // If there are pending updates
    if (lazy[node] !== 0) {

        // Updating the node
        tree[node] = (end - start + 1) * lazy[node];

        // Passing the update information to its children
        if (start !== end) {
            lazy[node * 2] = lazy[node];
            lazy[node * 2 + 1] = lazy[node];
        }

        // Resetting the lazy value for the current node
        lazy[node] = 0;
    }

    // At this point we are sure that pending lazy updates
    // are done for current node. So we can return value
    if (start >= l && end <= r)
        return tree[node];

    // If not completely in range but overlaps, recur for
    // children,
    let mid = Math.floor((start + end) / 2);
    let p1 = queryRange(node * 2, start, mid, l, r);
    let p2 = queryRange(node * 2 + 1, mid + 1, end, l, r);

    // Use the result of children calls to update this node
    return (p1 + p2);
}

// Hardcoded input
let n = 5, q = 3;
let queries = [[1, 0, 3, 5], [2, 1, 4], [1, 2, 4, 10]];

for (let i = 0; i < q; i++) {
    let type = queries[i][0];

    if (type === 1) {
        let l = queries[i][1], r = queries[i][2],
            v = queries[i][3];
        updateRange(1, 0, n - 1, l, r - 1, v);
    }

    else {
        let l = queries[i][1], r = queries[i][2];
        console.log(queryRange(1, 0, n - 1, l, r - 1));
    }
}


// This code is contributed by Sakshi

Output

10

Time Complexity: O(QlogN), where Q is the number of queries and N is the size of input array.
Auxiliary Space: O(N).