Skip to content
Related Articles

Related Articles

Improve Article

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

  • Difficulty Level : Expert
  • Last Updated : 03 Feb, 2021

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: 

C++




// 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;
}

Python3




# Python3 implementation of the approach
 
# Structure which contains
# all the necessary functions
# to solve a segment tree
class SegmentTemplate:
     
    def __init__(self) -> None:
         
        self.start = 0
        self.end = 0
        self.count = 0
        self.pendingUpdate = False
 
    # The function which
    # contains merge logic
    def merge(self, leftSub, rightSub) -> None:
 
        # If the left subarray
        # contains some pending update
        if (leftSub.pendingUpdate):
            self.count += (leftSub.end - leftSub.start +
                                     1 - leftSub.count)
        else:
            self.count += leftSub.count
 
        # If the right subarray
        # contains some pending update
        if (rightSub.pendingUpdate):
            self.count += (rightSub.end - rightSub.start +
                                      1 - rightSub.count)
        else:
            self.count += rightSub.count
 
    # To return the count
    def query(self) -> int:
 
        return self.count
 
    # If the segment tree
    # has pending update
    def hasPendingUpdate(self) -> bool:
 
        return self.pendingUpdate
 
    # Apply pending update
    def applyPendingUpdate(self) -> None:
 
        self.count = self.end - self.start + 1 - self.count
        self.pendingUpdate = False
 
    # For a node to
    # add pending update
    def addUpdate(self) -> None:
 
        self.pendingUpdate = not self.pendingUpdate
 
# The class that performs
# all the segment tree functions
class SegmentTree:
     
    def __init__(self, N) -> None:
         
        self.treeBuild = [SegmentTemplate() for _ in range(
            self._getSegmentTreeSize(N))]
        self.N = N
        self._buildTree(1, 0, N - 1)
 
    # For the query
    def query(self, start: int, end: int) -> int:
 
        # Stores the result
        result = self._query(1, start, end)
        return result.query()
 
    # Updates at the
    # specific location
    def update(self, start: int, end: int) -> None:
 
        self._update(1, start, end)
 
    def _buildTree(self, stIndex: int, start: int, end: int) -> None:
 
        # Defining index in treeBuild
        self.treeBuild[stIndex].start = start
        self.treeBuild[stIndex].end = end
 
        if (start == end):
            return
 
        # Dividing the array into two
        # parts for assigning the indices
        # and building segment tree
        mid = (start + end) // 2
        leftChildIndex = 2 * stIndex
        rightChildIndex = leftChildIndex + 1
 
        self._buildTree(leftChildIndex, start, mid)
        self._buildTree(rightChildIndex, mid + 1, end)
 
        # Combine the array
        self.treeBuild[stIndex].merge(self.treeBuild[leftChildIndex],
                                      self.treeBuild[rightChildIndex])
 
    # Gets the segment tree size
    def _getSegmentTreeSize(self, N: int) -> int:
 
        size = 1
        while size < N:
            size <<= 1
             
        return size << 1
 
    # Function for getting
    # to know the queries
    def _query(self, stIndex: int,
             start: int, end: int) -> SegmentTemplate:
 
        # When reaches a particular index
        if (self.treeBuild[stIndex].start == start and
            self.treeBuild[stIndex].end == end):
            result = self.treeBuild[stIndex]
            if (result.hasPendingUpdate()):
                result.applyPendingUpdate()
                 
            return result
 
        # Dividing the array into two
        # parts for assigning the indices
        # and querying segment trees
        mid = (self.treeBuild[stIndex].start +
               self.treeBuild[stIndex].end) // 2
                
        leftChildIndex = stIndex * 2
        rightChildIndex = leftChildIndex + 1
        result = SegmentTemplate()
 
        if (start > mid):
            result = self._query(rightChildIndex, start, end)
        elif (end <= mid):
            result = self._query(leftChildIndex, start, end)
        else:
            leftResult = self._query(leftChildIndex, start, mid)
            rightResult = self._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 (self.treeBuild[stIndex].hasPendingUpdate()):
            result.addUpdate()
            result.applyPendingUpdate()
 
        return result
 
    # Function for updating the tree
    def _update(self, stIndex: int, start: int, end: int) -> None:
 
        # Adds the update in O(1)
        # time to the node
        if (self.treeBuild[stIndex].start == start and
            self.treeBuild[stIndex].end == end):
            self.treeBuild[stIndex].addUpdate()
            return
 
        # Dividing the array into two
        # parts for assigning the indices
        # and querying segment trees
        mid = (self.treeBuild[stIndex].start +
               self.treeBuild[stIndex].end) // 2
                
        leftChildIndex = stIndex * 2
        rightChildIndex = leftChildIndex + 1
 
        if (start > mid):
            self._update(rightChildIndex, start, end)
        elif (end <= mid):
            self._update(leftChildIndex, start, end)
        else:
            self._update(leftChildIndex, start, mid)
            self._update(rightChildIndex, mid + 1, end)
 
        # Calling the build function
        # for building the segment tree
        self.treeBuild[stIndex].merge(self.treeBuild[leftChildIndex],
                                      self.treeBuild[rightChildIndex])
 
# Driver code
if __name__ == "__main__":
 
    N = 7
    st = SegmentTree(N)
 
    # Updating the segment tree
    st.update(2, 5)
 
    # Executing the query
    print("{}".format(st.query(1, 6)))
 
    # Updating the segment tree
    st.update(1, 3)
 
    # Executing the query
    print("{}".format(st.query(1, 6)))
 
# This code is contributed by sanjeev2552
Output: 
4
3

 




My Personal Notes arrow_drop_up
Recommended Articles
Page :