Range queries to count 1s in a subarray after flip operations

Given an array ‘arr’ all the numbers of which are initialized to ‘0’ then the array can be updated at any time in the following way:
All the elements of any sub-array from the array can be flipped i.e. ‘1’ => ‘0’ or ‘0’ => ‘1’.
The task is to find the number of 1s from the array for any incoming query [l, r] where ‘l’ and ‘r’ are valid indices in the given array.

Examples:

Input: n = 7
arr = {0, 0, 0, 0, 0, 0, 0}
Update from index '2' to '5', {0, 0, 1, 1, 1, 1, 0}
Query from index '1' to '6'
Update from index '1' to '3', {0, 1, 0, 0, 1, 1, 0}
Query from index '1' to '4'
Output:
4
2

Input: n = 5
arr = {0, 0, 0, 0, 0}
Update from index '0' to '2', {1, 1, 1, 0, 0}
Query from index '2' to '4'
Update from index '2' to '4', {1, 1, 0, 1, 1}
Query from index '0' to '3'
Output:
1
3


Approach: This problem can be solved using segment trees, you can read more about segment trees here.

  • Initialize the segment tree with value ‘0’ at all the nodes.
  • At any update, use update function to store updated values at that index only so that queries can be performed in O(Log n) time which is known as Lazy Propogation Method.
  • While building the segment tree, use merge function such that:
    • If pending updates are present in left and right subArrays, increment count by leftSub.end - leftSub.start + 1 - leftSub.count and rightSub.end - rightSub.start + 1 - rightSub.count respectively.
    • Otherwise, add leftSub.count and rightSub.count.

Below is the implementation of the above approach:

filter_none

edit
close

play_arrow

link
brightness_4
code

// C++ implementation of the approach
#include <bits/stdc++.h>
  
// Structure which contains
// all the necessary functions
// to solve a segment tree
struct SegmentTemplate {
    int start, end;
    int count = 0;
    bool pendingUpdate = false;
  
    // The function which
    // contains merge logic
    void merge(SegmentTemplate& leftSub, 
              SegmentTemplate& rightSub)
    {
        // If the left subarray
        // contains some pending update
        if (leftSub.pendingUpdate)
            count += leftSub.end - leftSub.start +
                                 1 - leftSub.count;
        else
            count += leftSub.count;
  
        // If the right subarray
        // contains some pending update
        if (rightSub.pendingUpdate)
            count += rightSub.end - rightSub.start +
                                 1 - rightSub.count;
        else
            count += rightSub.count;
    }
  
    // To return the count
    int query()
    {
        return count;
    }
  
    // If the segment tree
    // has pending update
    bool hasPendingUpdate()
    {
        return pendingUpdate;
    }
  
    // Apply pending update
    void applyPendingUpdate()
    {
        count = end - start + 1 - count;
        pendingUpdate = false;
    }
  
    // For a node to
    // add pending update
    void addUpdate()
    {
        pendingUpdate = !pendingUpdate;
    }
};
  
// The class that performs
// all the segment tree functions
class SegmentTree {
    SegmentTemplate* treeBuild;
    int N;
  
public:
    SegmentTree(int N)
    {
        this->N = N;
        treeBuild = new SegmentTemplate[getSegmentTreeSize(N)];
        buildTree(1, 0, N - 1);
    }
  
    // For the query
    int query(int start, int end)
    {
  
        // Stores the result
        SegmentTemplate result
            = query(1, start, end);
        return result.query();
    }
  
    // Updates at the
    // specific location
    void update(int start, int end)
    {
        update(1, start, end);
    }
  
private:
    void buildTree(int stIndex, int start, int end)
    {
  
        // Defining index in treeBuild
        treeBuild[stIndex].start
            = start,
            treeBuild[stIndex].end = end;
  
        if (start == end) {
            return;
        }
  
        // Dividing the array into two
        // parts for assigning the indices
        // and building segment tree
        int mid = (start + end) / 2,
            leftChildIndex = 2 * stIndex,
            rightChildIndex = leftChildIndex + 1;
  
        buildTree(leftChildIndex, start, mid);
        buildTree(rightChildIndex, mid + 1, end);
  
        // Combine the array
        treeBuild[stIndex]
            .merge(treeBuild[leftChildIndex],
                   treeBuild[rightChildIndex]);
    }
  
    // Gets the segment tree size
    int getSegmentTreeSize(int N)
    {
        int size = 1;
        for (; size < N; size <<= 1)
            ;
        return size << 1;
    }
  
    // Function for getting
    // to know the queries
    SegmentTemplate query(int stIndex, int start, int end)
    {
  
        // When reaches a particular index
        if (treeBuild[stIndex].start == start
            && treeBuild[stIndex].end == end) {
            SegmentTemplate result = treeBuild[stIndex];
            if (result.hasPendingUpdate())
                result.applyPendingUpdate();
            return result;
        }
  
        // Dividing the array into two
        // parts for assigning the indices
        // and querying segment trees
        int mid = (treeBuild[stIndex].start
                   + treeBuild[stIndex].end)
                  / 2,
            leftChildIndex = stIndex * 2,
            rightChildIndex = leftChildIndex + 1;
        SegmentTemplate result;
  
        if (start > mid)
            result = query(rightChildIndex, start, end);
        else if (end <= mid)
            result = query(leftChildIndex, start, end);
        else {
            SegmentTemplate leftResult = query(leftChildIndex, start, mid),
                        rightResult = query(rightChildIndex, mid + 1, end);
            result.start = leftResult.start;
            result.end = rightResult.end;
            result.merge(leftResult, rightResult);
        }
  
        // If tree has update,
        // apply the pending update
        if (treeBuild[stIndex].hasPendingUpdate()) {
            result.addUpdate();
            result.applyPendingUpdate();
        }
        return result;
    }
  
    // Function for updating the tree
    void update(int stIndex, int start, int end)
    {
  
        // Adds the update in O(1)
        // time to the node
        if (treeBuild[stIndex].start == start
            && treeBuild[stIndex].end == end) {
            treeBuild[stIndex].addUpdate();
            return;
        }
  
        // Dividing the array into two
        // parts for assigning the indices
        // and querying segment trees
        int mid = (treeBuild[stIndex].start
                   + treeBuild[stIndex].end)
                  / 2,
            leftChildIndex = stIndex * 2,
            rightChildIndex = leftChildIndex + 1;
  
        if (start > mid)
            update(rightChildIndex, start, end);
        else if (end <= mid)
            update(leftChildIndex, start, end);
        else {
            update(leftChildIndex, start, mid);
            update(rightChildIndex, mid + 1, end);
        }
  
        // Calling the build function
        // for building the segment tree
        treeBuild[stIndex]
            .merge(treeBuild[leftChildIndex],
                   treeBuild[rightChildIndex]);
    }
};
  
// Driver function
int main()
{
    int N = 7;
    SegmentTree st(N);
  
    // Updating the segment tree
    st.update(2, 5);
  
    // Executing the query
    printf("%d\n", st.query(1, 6));
  
    // Updating the segment tree
    st.update(1, 3);
  
    // Executing the query
    printf("%d\n", st.query(1, 6));
    return 0;
}

chevron_right


Output:

4
3


My Personal Notes arrow_drop_up

Coder only XD Do send reference if contacting me

If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.