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: 
 

  1. Add a value X to all the element from from index A to B where 0 ≤ A ≤ B ≤ N-1.
  2. 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 = 10 
Output: 
Sum of values in given range = 15 
Updated sum of values in given range = 45 
Explanation: 
Sum of values in the range 1 to 3 is 3 + 5 + 7 = 15. 
arr[] after adding 10 from index 1 to 5 is arr[] = {1, 13, 15, 17, 19, 21} 
Sum of values in the range 1 to 3 after update is 13 + 15 + 17 = 45.
Input: arr[] = { 11, 32, 5, 7, 19, 11, 8}, L = 2, R = 6, A = 1, B = 5, X = 16 
Output: 
Sum of values in given range = 50 
Updated sum of values in given range = 114 
Explanation: 
Sum of values in the range 2 to 6 is 5 + 7 + 19 + 11 + 8 = 50. 
arr[] after adding 16 from index 1 to 5 is arr[] = {11, 48, 21, 23, 35, 27, 8} 
Sum of values in the range 2 to 6 after update is 21 + 23 + 35 + 27 + 8 = 114. 
 

 

Approach: 
The recursive approach using a Segment Tree for the given problem is discussed in this article. In this post we will discussed an approach using Stack Data Structure to avoid recursion
Below are the steps to implement Segment Tree using Stack
 



  1. The idea is to use tuple to store the state which has Node number and range indexes in the Stack.
  2. For Building Segment Tree: 
    • Push the root Node to the stack as a tuple: 
       
Stack S;
start = 0, end = arr_size - 1
S.emplace(1, start, end)

mid = (start + end) / 2 
st.emplace(curr_node * 2, start, mid) 
st.emplace(curr_node * 2 + 1, mid + 1, end) 
 

  1.  
  2. If start index and end index is same as INF, then update the Segment Tree value at current index as: 
     

Value at current index is updated as value at left child and right child: 
tree[curr_node] = tree[2*curr_node] + tree[2*curr_node + 1] 
 

  1.  
  2. For Update Tree: 
    • Push the root node to the stack as done for building Segment Tree.
    • Pop the element from the stack and do the following untill stack becomes empty: 
      1. If the current node has any pending update then first update to the current node.
      2. If the current node ranges lies completely in the update query range, then update the current node with that value.
      3. If the current node ranges overlap with the update query range, then follow the above approach and push the tuple for left child and right child in the Stack.
      4. Update the query using the result of left and right child above.
  3. For Update Query: 
    • Push the root node to the stack as done for building Segment Tree.
    • Pop the element from the stack and do the following untill stack becomes empty: 
      1. If the current node ranges lies outside the given query, then continue with the next iteration.
      2. If the current node ranges lies completely in the update query range, then update the result with the current node value.
      3. If the current node ranges overlap with the update query range, then follow the above approach and push the tuple for left child and right child in the Stack.
      4. Update the result using the value obtained from left and right child Node.

Below is the implementation of the above approach:
 

filter_none

edit
close

play_arrow

link
brightness_4
code

#include"bits/stdc++.h"
using namespace std;
 
constexpr static int MAXSIZE = 1000;
constexpr static int INF
    = numeric_limits<int>::max();
 
// Segment Tree array
int64_t tree[MAXSIZE];
 
// Lazy Update array
int64_t lazy[MAXSIZE];
 
// This tuple will hold tree state
// the stacks
using QueryAdaptor
    = tuple<int64_t,
            int64_t,
            int64_t>;
 
// Build our segment tree
void build_tree(int64_t* arr,
                int64_t arr_size)
{
   
    // Stack will use to update
    // the tree value
    stack<QueryAdaptor> st;
 
    // Emplace the root of the tree
    st.emplace(1, 0, arr_size - 1);
 
    // Repeat until empty
    while (!st.empty()) {
 
        // Take the indexes at the
        // top of the stack
        int64_t currnode, curra, currb;
 
        // value at the top of the
        // stack
        tie(currnode, curra, currb) = st.top();
 
        // Pop the value from the
        // stack
        st.pop();
 
        // Flag with INF ranges are merged
        if (curra == INF && currb == INF) {
            tree[currnode] = tree[currnode * 2]
                             + tree[currnode * 2 + 1];
        }
 
        // Leaf node
        else if (curra == currb) {
            tree[currnode] = arr[curra];
        }
 
        else {
 
            // Insert flag node inverse
            // order of evaluation
            st.emplace(currnode, INF, INF);
            int64_t mid = (curra + currb) / 2;
 
            // Push children
            st.emplace(currnode * 2,
                       curra, mid);
            st.emplace(currnode * 2 + 1,
                       mid + 1, currb);
        }
    }
}
 
// A utility function that propagates
// updates lazily down the tree
inline void push_down(int64_t node,
                      int64_t a,
                      int64_t b)
{
    if (lazy[node] != 0) {
        tree[node] += lazy[node] * (b - a + 1);
 
        if (a != b) {
            lazy[2 * node] += lazy[node];
            lazy[2 * node + 1] += lazy[node];
        }
 
        lazy[node] = 0;
    }
}
 
// Iterative Range_Update function to
// add val to all elements in the
// range i-j (inclusive)
void update_tree(int64_t arr_size,
                 int64_t i,
                 int64_t j,
                 int64_t val)
{
 
    // Intialize the stack
    stack<QueryAdaptor> st;
 
    // Emplace the root of the tree
    st.emplace(1, 0, arr_size - 1);
 
    // Work until empty
    while (!st.empty()) {
 
        // Take the indexes at the
        // top of the stack
        int64_t currnode, curra, currb;
        tie(currnode, curra, currb) = st.top();
        st.pop();
 
        // Flag with INF ranges are merged
        if (curra == INF && currb == INF) {
            tree[currnode] = tree[currnode * 2]
                             + tree[currnode * 2 + 1];
        }
 
        // Traverse the previous updates
        // down the tree
        else {
            push_down(currnode, curra, currb);
 
            // No overlap condition
            if (curra > currb || curra > j
                || currb < i) {
                continue;
            }
 
            // Total overlap condition
            else if (curra >= i && currb <= j)
            {
               
                // Update lazy array
                tree[currnode] += val * (currb - curra + 1);
 
                if (curra != currb) {
                    lazy[currnode * 2] += val;
                    lazy[currnode * 2 + 1] += val;
                }
            }
 
            // Partial Overlap
            else
            {
               
                // Insert flag node inverse
                // order of evaluation
                st.emplace(currnode, INF, INF);
 
                int64_t mid = (curra + currb) / 2;
 
                // Push children
                st.emplace(currnode * 2,
                           curra, mid);
                st.emplace(currnode * 2 + 1,
                           mid + 1, currb);
            }
        }
    }
}
 
// A function that find the sum of
// all elements in the range i-j
int64_t query(int64_t arr_size,
              int64_t i,
              int64_t j)
{
   
    // Intialize stack
    stack<QueryAdaptor> st;
 
    // Emplace root of the tree
    st.emplace(1, 0, arr_size - 1);
 
    int64_t result = 0;
 
    while (!st.empty())
    {
 
        // Take the indexes at the
        // top of the stack
        int64_t currnode, curra, currb;
        tie(currnode, curra, currb) = st.top();
        st.pop();
 
        // Traverse the previous updates
        // down the tree
        push_down(currnode, curra, currb);
 
        // No overlap
        if (curra > currb || curra > j
            || currb < i) {
            continue;
        }
 
        // Total Overlap
        else if (curra >= i && currb <= j) {
            result += tree[currnode];
        }
 
        // Partial Overlap
        else {
            std::int64_t mid = (curra + currb) / 2;
 
            // Push children
            st.emplace(currnode * 2,
                       curra, mid);
            st.emplace(currnode * 2 + 1,
                       mid + 1, currb);
        }
    }
 
    return result;
}
 
// Driver Code
int main()
{
   
    // Initialize our trees with 0
    memset(tree, 0, sizeof(int64_t) * MAXSIZE);
    memset(lazy, 0, sizeof(int64_t) * MAXSIZE);
 
    int64_t arr[] = { 1, 3, 5, 7, 9, 11 };
    int n = sizeof(arr) / sizeof(arr[0]);
 
    // Build segment tree from given array
    build_tree(arr, n);
 
    // Print sum of values in array
    // from index 1 to 3
    cout << "Sum of values in given range = "
         << query(n, 1, 3)
         << endl;
 
    // Add 10 to all nodes at indexes
    // from 1 to 5
    update_tree(n, 1, 5, 10);
 
    // Find sum after the value is updated
    cout << "Updated sum of values in given range = "
         << query(n, 1, 3)
         << endl;
 
    return 0;
}
chevron_right

filter_none

edit
close

play_arrow

link
brightness_4
code

// Java implementation of the approach
import java.util.Arrays;
import java.util.List;
import java.util.Stack;
class GFG
{
 
    static final int MAXSIZE = 1000;
    static final int INF = (int) Double.POSITIVE_INFINITY;
 
    // Segment Tree array
    static int[] tree = new int[MAXSIZE];
 
    // Lazy Update array
    static int[] lazy = new int[MAXSIZE];
 
    // Build our segment tree
    static void build_tree(int[] arr, int arr_size)
    {
       
        // Stack will use to update
        // the tree value
        Stack<List<Integer>> st = new Stack<>();
 
        // push the root of the tree
        st.push(Arrays.asList(1, 0, arr_size - 1));
 
        // Repeat until empty
        while (!st.isEmpty())
        {
 
            // Take the indexes at the
            // top of the stack
            int currnode, curra, currb;
 
            // value at the top of the
            // stack
            List<Integer> temp = st.peek();
            currnode = temp.get(0);
            curra = temp.get(1);
            currb = temp.get(2);
 
            // Pop the value from the
            // stack
            st.pop();
 
            // Flag with INF ranges are merged
            if (curra == INF && currb == INF)
            {
                tree[currnode] = tree[currnode * 2] +
                  tree[currnode * 2 + 1];
            }
 
            // Leaf node
            else if (curra == currb)
            {
                tree[currnode] = arr[curra];
            }
 
            else {
 
                // Insert flag node inverse
                // order of evaluation
                st.push(Arrays.asList(currnode, INF, INF));
 
                int mid = (curra + currb) / 2;
 
                // Push children
                st.push(Arrays.asList(currnode * 2, curra, mid));
                st.push(Arrays.asList(currnode * 2 + 1, mid + 1, currb));
            }
        }
    }
 
    // A utility function that propagates
    // updates lazily down the tree
    static void push_down(int node, int a, int b)
    {
        if (lazy[node] != 0)
        {
            tree[node] += lazy[node] * (b - a + 1);
 
            if (a != b)
            {
                lazy[2 * node] += lazy[node];
                lazy[2 * node + 1] += lazy[node];
            }
 
            lazy[node] = 0;
        }
    }
 
    // Iterative Range_Update function to
    // add val to all elements in the
    // range i-j (inclusive)
    static void update_tree(int arr_size, int i,
                            int j, int val)
    {
 
        // Intialize the stack
        Stack<List<Integer>> st = new Stack<>();
 
        // push the root of the tree
        st.push(Arrays.asList(1, 0, arr_size - 1));
 
        // Work until empty
        while (!st.isEmpty())
        {
 
            // Take the indexes at the
            // top of the stack
            int currnode, curra, currb;
            List<Integer> temp = st.peek();
            currnode = temp.get(0);
            curra = temp.get(1);
            currb = temp.get(2);
            st.pop();
 
            // Flag with INF ranges are merged
            if (curra == INF && currb == INF)
            {
                tree[currnode] = tree[currnode * 2] +
                  tree[currnode * 2 + 1];
            }
 
            // Traverse the previous updates
            // down the tree
            else
            {
                push_down(currnode, curra, currb);
 
                // No overlap condition
                if (curra > currb || curra > j || currb < i)
                {
                    continue;
                }
 
                // Total overlap condition
                else if (curra >= i && currb <= j)
                {
                   
                    // Update lazy array
                    tree[currnode] += val * (currb - curra + 1);
 
                    if (curra != currb)
                    {
                        lazy[currnode * 2] += val;
                        lazy[currnode * 2 + 1] += val;
                    }
                }
 
                // Partial Overlap
                else
                {
                   
                    // Insert flag node inverse
                    // order of evaluation
                    st.push(Arrays.asList(currnode, INF, INF));
 
                    int mid = (curra + currb) / 2;
 
                    // Push children
                    st.push(Arrays.asList(currnode * 2, curra, mid));
                    st.push(Arrays.asList(currnode * 2 + 1,
                                          mid + 1, currb));
                }
            }
        }
    }
 
    // A function that find the sum of
    // all elements in the range i-j
    static int query(int arr_size, int i, int j)
    {
       
        // Intialize stack
        Stack<List<Integer>> st = new Stack<>();
 
        // push root of the tree
        st.push(Arrays.asList(1, 0, arr_size - 1));
 
        int result = 0;
 
        while (!st.isEmpty())
        {
 
            // Take the indexes at the
            // top of the stack
            int currnode, curra, currb;
            List<Integer> temp = st.peek();
            currnode = temp.get(0);
            curra = temp.get(1);
            currb = temp.get(2);
            st.pop();
 
            // Traverse the previous updates
            // down the tree
            push_down(currnode, curra, currb);
 
            // No overlap
            if (curra > currb || curra > j || currb < i)
            {
                continue;
            }
 
            // Total Overlap
            else if (curra >= i && currb <= j)
            {
                result += tree[currnode];
            }
 
            // Partial Overlap
            else
            {
                int mid = (curra + currb) / 2;
 
                // Push children
                st.push(Arrays.asList(currnode * 2, curra, mid));
                st.push(Arrays.asList(currnode * 2 + 1, mid + 1, currb));
            }
        }
 
        return result;
    }
 
    // Driver Code
    public static void main(String[] args)
    {
 
        // Initialize our trees with 0
        Arrays.fill(tree, 0);
        Arrays.fill(lazy, 0);
 
        int arr[] = { 1, 3, 5, 7, 9, 11 };
        int n = arr.length;
 
        // Build segment tree from given array
        build_tree(arr, n);
 
        // Print sum of values in array
        // from index 1 to 3
        System.out.printf("Sum of values in given range = %d\n", query(n, 1, 3));
 
        // Add 10 to all nodes at indexes
        // from 1 to 5
        update_tree(n, 1, 5, 10);
 
        // Find sum after the value is updated
        System.out.printf("Updated sum of values in given range = %d\n", query(n, 1, 3));
    }
}
 
// This code is contributed by sanjeev2552
chevron_right

Output: 
Sum of values in given range = 15
Updated sum of values in given range = 45

 

Time Complexity: 
 

 

Attention reader! Don’t stop learning now. Get hold of all the important DSA concepts with the DSA Self Paced Course at a student-friendly price and become industry ready.




Article Tags :
Practice Tags :