Introduction to Segment Trees - Data Structure and Algorithm Tutorials

// Java code for segment tree with sum
// range and update query
import java.io.*;
import java.util.*;

class GFG {
    static int n = 6;
    static int A[] = { 0, 1, 3, 5, -2, 3 };

    // Create a segment tree of size 4*n
    static int ST[] = new int[4 * n];
    public static void build(int node, int L, int R)
    {

        // Leaf node where L == R
        if (L == R) {
            ST[node] = A[L];
        }
        else {

            // Find the middle element to
            // split the array into two halves
            int mid = (L + R) / 2;

            // Recursively travel the
            // left half
            build(2 * node, L, mid);

            // Recursively travel the
            // right half
            build(2 * node + 1, mid + 1, R);

            // Storing the sum of both the
            // children into the parent
            ST[node] = ST[2 * node] + ST[2 * node + 1];
        }
    }

    public static void update(int node, int L, int R,
                              int idx, int val)
    {

        // Find the lead node and
        // update its value
        if (L == R) {
            A[idx] += val;
            ST[node] += val;
        }
        else {

            // Find the mid
            int mid = (L + R) / 2;

            // If node value idx is at the
            // left part then update
            // the left part
            if (L <= idx && idx <= mid)
                update(2 * node, L, mid, idx, val);
            else
                update(2 * node + 1, mid + 1, R, idx, val);

            // Store the information in parents
            ST[node] = ST[2 * node] + ST[2 * node + 1];
        }
    }

    public static int query(int node, int tl, int tr, int l,
                            int r)
    {

        // If it lies out of range then
        // return 0
        if (r < tl || tr < l)
            return 0;

        // If the node contains the range then
        // return the node value
        if (l <= tl && tr <= r)
            return ST[node];
        int tm = (tl + tr) / 2;

        // Recursively traverse left and right
        // and find the node
        return query(2 * node, tl, tm, l, r)
            + query(2 * node + 1, tm + 1, tr, l, r);
    }

    // Driver Code
    public static void main(String[] args)
    {
        // Build a segment tree
        build(1, 0, n - 1);
        System.out.println(
            "Sum of values in range 0-4 are: "
            + query(1, 0, n - 1, 0, 4));

        // Update the value at idx = 1 by
        // 100 ths becoming 101
        update(1, 0, n - 1, 1, 100);
        System.out.println(
            "Value at index 1 increased by 100");
        System.out.println("sum of value in range 1-3 are: "
                           + query(1, 0, n - 1, 1, 3));
    }
}

// This code is contributed by Rohit Pradhan

// Java program to show segment tree to
// demonstrate lazy propagation
import java.util.*;

class GFG {

    // Ideally, we should not use global
    // variables and large constant-sized
    // arrays, we have done it here for
    // simplicity.

    static final int MAX = 1000;

    // To store segment tree
    static int[] tree = new int[MAX];

    // To store pending updates
    static int[] lazy = new int[MAX];

    // si -> index of current node in segment tree
    // ss and se -> Starting and ending
    // indices of elements for which current
    // nodes stores sum.
    // us and ue -> starting and ending indexes
    // of update query
    // diff -> which we need to add in the
    // range us to ue
    static void updateRangeUtil(int si, int ss, int se,
                                int us, int ue, int diff)
    {

        // If lazy value is non-zero for
        // current node of segment tree, then
        // there are some pending updates. So,
        // we need to make sure that the
        // pending updates are done before
        // making new updates. Because this
        // value may be used by parent after
        // recursive calls (See last line
        // of this function)
        if (lazy[si] != 0) {

            // Make pending updates using
            // value stored in lazy nodes
            tree[si] += (se - ss + 1) * lazy[si];

            // checking if it is not leaf node
            // because if it is leaf node then
            // we cannot go further
            if (ss != se) {

                // We can postpone updating
                // children we don't need
                // their new values now. Since
                // we are not yet updating
                // children of si, we need to
                // set lazy flags for the children
                lazy[si * 2 + 1] += lazy[si];
                lazy[si * 2 + 2] += lazy[si];
            }

            // Set the lazy value for current
            // node as 0 as it has been updated
            lazy[si] = 0;
        }

        // out of range
        if (ss > se || ss > ue || se < us) {
            return;
        }

        // Current segment is fully in range
        if (ss >= us && se <= ue) {

            // Add the difference to
            // current node
            tree[si] += (se - ss + 1) * diff;

            // Same logic for checking
            // leaf node or not
            if (ss != se) {

                // This is where we store
                // values in lazy nodes,
                // rather than updating the
                // segment tree itself Since
                // we don't need these updated
                // values now we postpone
                // updates by storing values
                // in lazy[]
                lazy[si * 2 + 1] += diff;
                lazy[si * 2 + 2] += diff;
            }
            return;
        }

        // If not completely in rang,
        // but overlaps, recur for children,
        int mid = (ss + se) / 2;
        updateRangeUtil(si * 2 + 1, ss, mid, us, ue, diff);
        updateRangeUtil(si * 2 + 2, mid + 1, se, us, ue,
                        diff);

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

    // Function to update a range of values
    // in segment tree
    // us and eu -> starting and ending
    // indices of update query
    // ue -> ending index of update query
    // diff -> which we need to add in the
    // range us to ue
    static void updateRange(int n, int us, int ue, int diff)
    {
        updateRangeUtil(0, 0, n - 1, us, ue, diff);
    }

    // A recursive function to get the sum of
    // values in given range of the array.
    // The following are parameters for this function.
    // si --> Index of current node in the st.
    // Initially 0 is passed as root is always
    // at' index 0
    // ss & se --> Starting and ending indices
    // of the segment represented by current
    // node, i.e., tree[si]
    // qs & qe --> Starting and ending
    // indices of query range
    static int getSumUtil(int ss, int se, int qs, int qe,
                          int si)
    {

        // If lazy flag is set for current
        // node of segment tree, then there
        // are some pending updates. So we
        // need to make sure that the pending
        // updates are done before
        // processing the sub sum query
        if (lazy[si] != 0) {

            // Make pending updates to this
            // node. Note that this node
            // represents sum of elements in
            // arr[ss..se]  and all these
            // elements must be increased by
            // lazy[si]
            tree[si] += (se - ss + 1) * lazy[si];

            // Checking if it is not leaf node
            // because if it is leaf node then
            // we cannot go further
            if (ss != se) {

                // Since we are not yet
                // updating children os si,
                // we need to set lazy values
                // for the children
                lazy[si * 2 + 1] += lazy[si];
                lazy[si * 2 + 2] += lazy[si];
            }

            // unset the lazy value for current
            // node as it has been updated
            lazy[si] = 0;
        }

        // Out of range
        if (ss > se || ss > qe || se < qs) {
            return 0;
        }

        // At this point we are sure that
        // pending lazy updates are done for
        // current node. So we can return
        // value

        // If this segment lies in range
        if (ss >= qs && se <= qe) {
            return tree[si];
        }

        // If a part of this segment overlaps
        // with the given range
        int mid = (ss + se) / 2;
        return getSumUtil(ss, mid, qs, qe, 2 * si + 1)
            + getSumUtil(mid + 1, se, qs, qe, 2 * si + 2);
    }

    // Return sum of elements in range from
    // index qs (query start) to qe (query end).
    // It mainly uses getSumUtil()
    static int getSum(int n, int qs, int qe)
    {

        // Check for erroneous input values
        if (qs < 0 || qe > n - 1 || qs > qe) {
            System.out.println("Invalid Input");
            return -1;
        }
        return getSumUtil(0, n - 1, qs, qe, 0);
    }

    // A recursive function that constructs
    // Segment Tree for array[ss..se].
    // si is index of current node in st.
    static void constructSTUtil(int arr[], int ss, int se,
                                int si)
    {

        // Out of range as ss can never
        // be greater than se
        if (ss > se)
            return;

        // If there is one element in array,
        // store it in current node of segment
        // tree and return
        if (ss == se) {
            tree[si] = arr[ss];
            return;
        }

        // If there are more than one elements,
        // then recur for left and right
        // subtrees and store the sum of
        // values in this node
        int mid = (ss + se) / 2;
        constructSTUtil(arr, ss, mid, si * 2 + 1);
        constructSTUtil(arr, mid + 1, se, si * 2 + 2);

        tree[si] = tree[si * 2 + 1] + tree[si * 2 + 2];
    }

    // Function to construct segment tree
    // from given array.This function allocates
    // memory for segment tree and calls
    // constructSTUtil() to fill the
    // allocated memory
    static void constructST(int arr[], int n)
    {

        // Fill the allocated memory st
        constructSTUtil(arr, 0, n - 1, 0);
    }

    // Driver program to test above functions
    public static void main(String[] args)
    {
        int arr[] = { 1, 3, 5, 7, 9, 11 };
        int n = arr.length;

        // Build segment tree from given array
        constructST(arr, n);

        // Print sum of values in array
        // from index 1 to 3
        System.out.println("Sum of values in given range = "
                           + getSum(n, 1, 3));

        // Add 10 to all nodes at indexes
        // from 1 to 5.
        updateRange(n, 1, 5, 10);

        // Find sum after the value is updated
        System.out.println(
            "Updated sum of values in given range = "
            + getSum(n, 1, 3));
    }
}
// This code is contributed by Prasad Kandekar(prasad264)