Open In App

Segment Tree for Range Assignment and Range Sum Queries

Last Updated : 20 Feb, 2024
Improve
Improve
Like Article
Like
Save
Share
Report

Given an array of size filled with all 0s, the task is to answer 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).



Similar Reads

Segment Tree for Range Bitwise OR and Range Minimum Queries
There is an array of n elements, initially filled with zeros. The task is to perform q queries from given queries[][] and there are two types of queries, type 1 queries in the format {1, l, r, v} and type 2 queries in the format {2, l, r}. Type 1: Apply the operation ai = ai | v (bitwise OR) to all elements on the segment l to r.Type 2: Find the mi
21 min read
Queries for elements having values within the range A to B in the given index range using Segment Tree
Given an array arr[] of N elements and two integers A to B, the task is to answer Q queries each having two integers L and R. For each query, the task is to find the number of elements in the subarray arr[L...R] which lies within the range A to B (both included). Examples: Input: arr[] = {7, 3, 9, 13, 5, 4}, A=4, B=7 query = { 1, 5 } Output: 2 Expl
13 min read
Segment Tree for Range Multiplication and Sum Queries
Given an array of size N filled with all 1s, the task is to answer Q queries, where the queries can be one of the two types: Type 1 (1, L, R, X): Multiply all elements on the segment from L to R−1 by X, and Type 2 (2, L, R): Find the sum on the segment from L to R−1. Examples: Input: N = 5, Q = 6, queries[] = { {1, 0, 3, 3}, {2, 1, 2}, {1, 1, 4, 4}
19 min read
Queries for greatest pair sum in the given index range using Segment Tree
Given an array arr[] containing N integers and an array Q[] representing the range [L, R], the task is to find the greatest pair sum value in the range [L, R] where 0 ? L ? R ? N - 1. Examples: Input: arr[] = {1, 3, 2, 7, 9, 11}, Q[] = {1, 4} Output: 16 Explanation: The greatest pair sum in range 1 to 4 is 7 + 9 = 16. Input: arr[] = {0, 1, 2, 3, 4,
14 min read
Queries for the count of even digit sum elements in the given range using Segment Tree.
Given an array arr[] of N elements, the task is to answer Q queries each having two integers L and R. For each query, the task is to find the number of elements in the subarray arr[L...R] whose digit sum is even.Examples: Input: arr[] = {7, 3, 19, 13, 5, 4} query = { 1, 5 } Output: 3 Explanation: Elements 19, 13 and 4 have even digit sum in the sub
16 min read
Queries for elements greater than K in the given index range using Segment Tree
Given an array arr[] of N elements and a number of queries where each query will contain three integers L, R, and K. For each query, the task is to find the number of elements in the subarray arr[L...R] which are greater than K. Examples: Input: arr[] = {7, 3, 9, 13, 5, 4}, q[] = {{0, 3, 6}, {1, 5, 8}} Output: 3 2 Query 1: Only 7, 9 and 13 are grea
26 min read
Two equal sum segment range queries
Given an array arr[] consisting of N positive integers, and some queries consisting of a range [L, R], the task is to find whether the sub-array from the given index range can be divided into two contiguous parts of non-zero length and equal sum.Examples: Input: arr[] = {1, 1, 2, 3}, q[] = {{0, 1}, {1, 3}, {1, 2}} Output: Yes Yes No q[0]: The sub-a
9 min read
Dynamic Segment Trees : Online Queries for Range Sum with Point Updates
Prerequisites: Segment TreeGiven a number N which represents the size of the array initialized to 0 and Q queries to process where there are two types of queries: 1 P V: Put the value V at position P.2 L R: Output the sum of values from L to R. The task is to answer these queries. Constraints: 1 ? N ? 1018Q ? 1051 ? L ? R? N Note: Queries are onlin
22 min read
Segment Tree for Range Addition and Range Minimum Query
Given an array of size N initially filled with 0s. There are Q queries to be performed, and each query can be one of two types: Type 1 (1, L, R, X): Add X to all elements in the segment from L to R−1.Type 2 (2, L, R): Find the minimum value in the segment from L to R−1.Example: Input: n = 5, q = 6, queries = { {1, 0, 3, 3}, {2, 1, 2}, {1, 1, 4, 4},
16 min read
Range Sum and Update in Array : Segment Tree using Stack
Given an array arr[] of N integers. The task is to do the following operations: Add a value X to all the element from index A to B where 0 ? A ? B ? N-1.Find the sum of the element from index L to R where 0 ? L ? R ? N-1 before and after the update given to the array above.Example: Input: arr[] = {1, 3, 5, 7, 9, 11}, L = 1, R = 3, A = 1, B = 5, X =
21 min read