Segment Tree for Range Assignment and Range Sum Queries

Last Updated : 20 Feb, 2024

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 ` `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).

Article Tags :
Practice Tags :