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 |
10
Time Complexity: O(QlogN), where Q is the number of queries and N is the size of input array.
Auxiliary Space: O(N).
Contact Us