Open In App

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

Improve
Improve
Like Article
Like
Save
Share
Report

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


Java




// Java implementation for the above approach.
 
import java.util.*;
 
// Class for defining the Segment Tree template
class SegmentTemplate {
   
    // Initialize start, end and count to 0, and pendingUpdate flag to false
    int start, end;
    int count = 0;
    boolean pendingUpdate = false;
 
    // Merge function to update the count value of the node
    void merge(SegmentTemplate leftSub, SegmentTemplate rightSub) {
 
        // If leftSub has a pending update, update the count with the new value
        // Else, update the count with leftSub's count value
        if (leftSub.pendingUpdate)
            count += leftSub.end - leftSub.start + 1 - leftSub.count;
        else
            count += leftSub.count;
 
        // If rightSub has a pending update, update the count with the new value
        // Else, update the count with rightSub's count value
        if (rightSub.pendingUpdate)
            count += rightSub.end - rightSub.start + 1 - rightSub.count;
        else
            count += rightSub.count;
    }
 
    // Query function to return the count value of the node
    int query() {
        return count;
    }
 
    // Check if the node has a pending update
    boolean hasPendingUpdate() {
        return pendingUpdate;
    }
 
    // Apply the pending update to the node
    void applyPendingUpdate() {
        count = end - start + 1 - count;
        pendingUpdate = false;
    }
 
    // Function to add a new update to the node
    void addUpdate() {
        pendingUpdate = !pendingUpdate;
    }
 
}
 
class SegmentTree {
   
    SegmentTemplate[] treeBuild;
    int N;
 
      // Constructor that takes an integer N as input
    public SegmentTree(int N) {
       
      // Assigning the input value to a class variable
      this.N = N;
       
      // Creating an array of SegmentTemplate objects based on
      // the size of the segment tree
      treeBuild = new SegmentTemplate[getSegmentTreeSize(N)];
       
      // Building the segment tree
      buildTree(1, 0, N - 1);
    }
     
      // Method to query the segment tree for a given range
    int query(int start, int end) {
       
          // Calling the query method of the SegmentTemplate class with the
          // given range and storing the result in a SegmentTemplate object
        SegmentTemplate result = query(1, start, end);
       
          // Returning the result of the query
        return result.query();
    }
 
      // Method to update the segment tree for a given range
    void update(int start, int end) {
       
          // Calling the update method of the SegmentTemplate class with the given range
        update(1, start, end);
    }
 
      // Method to build the segment tree recursively
    void buildTree(int stIndex, int start, int end) {
       
          // Creating a new SegmentTemplate object at the
          // current index of the treeBuild array
        treeBuild[stIndex] = new SegmentTemplate();
       
          // Setting the start index of the current SegmentTemplate object
        treeBuild[stIndex].start = start;
       
          // Setting the end index of the current SegmentTemplate object
        treeBuild[stIndex].end = end;
 
          // Base case: if start and end are equal, return
        if (start == end) {
            return;
        }
 
          // Calculating the midpoint index of the current range
        int mid = (start + end) / 2;
       
          // Calculating the index of the left child of the current node
        int leftChildIndex = 2 * stIndex;
       
          // Calculating the index of the right child of the current node
        int rightChildIndex = leftChildIndex + 1;
 
          // Recursively building the left subtree
        buildTree(leftChildIndex, start, mid);
       
          // Recursively building the right subtree
        buildTree(rightChildIndex, mid + 1, end);
 
        // Merging the values of the left and right subtrees into the current node         
        treeBuild[stIndex].merge(treeBuild[leftChildIndex], treeBuild[rightChildIndex]);
    }
 
      // Method to calculate the size of the segment tree based on the input size N
    int getSegmentTreeSize(int N) {
       
          // Starting with a size of 1
        int size = 1;
       
          // Doubling the size until it is greater than or equal to N
        while (size < N) {
            size <<= 1;
        }
       
          // Doubling the size again and returning the final result
        return size << 1;
    }
 
      // Method to query the segment tree for a given range
    SegmentTemplate query(int stIndex, int start, int end) {
       
        // Base case: if the current node represents the given range, return the node
        if (treeBuild[stIndex].start == start && treeBuild[stIndex].end == end) {
            SegmentTemplate result = treeBuild[stIndex];
 
              // If the node has a pending update, apply it before returning the node
              if (result.hasPendingUpdate())
                  result.applyPendingUpdate();
 
            return result;
        }
 
          // Calculating the midpoint index of the current range
        int mid = (treeBuild[stIndex].start + treeBuild[stIndex].end) / 2;
       
          // Calculating the index of the left child of the current node
        int leftChildIndex = stIndex * 2;
           
          // Calculating the index of the right child of the current node
        int rightChildIndex = leftChildIndex + 1;
        SegmentTemplate result;
 
          // If the start index of the query range is greater than the midpoint,
          // recurse on the right subtree
        if (start > mid)
            result = query(rightChildIndex, start, end);
       
          // If the end index of the query range is less than or equal to the midpoint,
          // recurse on the left subtree
        else if (end <= mid)
            result = query(leftChildIndex, start, end);
       
          // Otherwise, recurse on both subtrees and merge the results
        else {
            SegmentTemplate leftResult = query(leftChildIndex, start, mid);
            SegmentTemplate rightResult = query(rightChildIndex, mid + 1, end);
            result = new SegmentTemplate();
            result.start = leftResult.start;
            result.end = rightResult.end;
            result.merge(leftResult, rightResult);
        }
 
          // If the current node has a pending update,
          // apply it to the result before returning
        if (treeBuild[stIndex].hasPendingUpdate()) {
            result.addUpdate();
            result.applyPendingUpdate();
        }
       
        return result; // Return the result of the query
    }
 
 
    // 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;
        int leftChildIndex = stIndex * 2;
        int 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
    public static void main(String[] args)
    {
        int N = 7;
        SegmentTree st = new SegmentTree(N);
 
        // Updating the segment tree
        st.update(2, 5);
 
        // Executing the query
        System.out.println(st.query(1, 6));
 
        // Updating the segment tree
        st.update(1, 3);
 
        // Executing the query
        System.out.println(st.query(1, 6));
 
    }
};
 
// This code is contributed by amit_mangal_


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


Javascript




class SegmentTemplate {
    constructor() {
        this.start = 0;
        this.end = 0;
        this.count = 0;
        this.pendingUpdate = false;
    }
 
    merge(leftSub, rightSub) {
        if (leftSub.pendingUpdate)
            this.count += leftSub.end - leftSub.start + 1 - leftSub.count;
        else
            this.count += leftSub.count;
 
        if (rightSub.pendingUpdate)
            this.count += rightSub.end - rightSub.start + 1 - rightSub.count;
        else
            this.count += rightSub.count;
    }
 
    query() {
        return this.count;
    }
 
    hasPendingUpdate() {
        return this.pendingUpdate;
    }
 
    applyPendingUpdate() {
        this.count = this.end - this.start + 1 - this.count;
        this.pendingUpdate = false;
    }
 
    addUpdate() {
        this.pendingUpdate = !this.pendingUpdate;
    }
}
 
class SegmentTree {
    constructor(N) {
        this.N = N;
        this.treeBuild = new Array(this.getSegmentTreeSize(N));
        this.buildTree(1, 0, N - 1);
    }
 
    query(start, end) {
        const queryResult = this.querySegment(1, start, end);
        return queryResult.query();
    }
 
    update(start, end) {
        this.updateSegment(1, start, end);
    }
 
    buildTree(stIndex, start, end) {
        this.treeBuild[stIndex] = new SegmentTemplate();
        this.treeBuild[stIndex].start = start;
        this.treeBuild[stIndex].end = end;
 
        if (start === end) {
            return;
        }
 
        const mid = Math.floor((start + end) / 2);
        const leftChildIndex = 2 * stIndex;
        const rightChildIndex = leftChildIndex + 1;
 
        this.buildTree(leftChildIndex, start, mid);
        this.buildTree(rightChildIndex, mid + 1, end);
 
        this.treeBuild[stIndex].merge(this.treeBuild[leftChildIndex],
        this.treeBuild[rightChildIndex]);
    }
 
    getSegmentTreeSize(N) {
        let size = 1;
        while (size < N) {
            size <<= 1;
        }
        return size << 1;
    }
 
    querySegment(stIndex, start, end) {
        if (this.treeBuild[stIndex].start === start && this.treeBuild[stIndex].end === end) {
            const result = this.treeBuild[stIndex];
            if (result.hasPendingUpdate())
                result.applyPendingUpdate();
            return result;
        }
 
        const mid = Math.floor((this.treeBuild[stIndex].start + this.treeBuild[stIndex].end) / 2);
        const leftChildIndex = stIndex * 2;
        const rightChildIndex = leftChildIndex + 1;
 
        let queryResult;
 
        if (start > mid)
            queryResult = this.querySegment(rightChildIndex, start, end);
        else if (end <= mid)
            queryResult = this.querySegment(leftChildIndex, start, end);
        else {
            const leftResult = this.querySegment(leftChildIndex, start, mid);
            const rightResult = this.querySegment(rightChildIndex, mid + 1, end);
            queryResult = new SegmentTemplate();
            queryResult.start = leftResult.start;
            queryResult.end = rightResult.end;
            queryResult.merge(leftResult, rightResult);
        }
 
        if (this.treeBuild[stIndex].hasPendingUpdate()) {
            queryResult.addUpdate();
            queryResult.applyPendingUpdate();
        }
 
        return queryResult;
    }
 
    updateSegment(stIndex, start, end) {
        if (this.treeBuild[stIndex].start === start && this.treeBuild[stIndex].end === end) {
            this.treeBuild[stIndex].addUpdate();
            return;
        }
 
        const mid = Math.floor((this.treeBuild[stIndex].start + this.treeBuild[stIndex].end) / 2);
        const leftChildIndex = stIndex * 2;
        const rightChildIndex = leftChildIndex + 1;
 
        if (start > mid) {
            this.updateSegment(rightChildIndex, start, end);
        } else if (end <= mid) {
            this.updateSegment(leftChildIndex, start, end);
        } else {
            this.updateSegment(leftChildIndex, start, mid);
            this.updateSegment(rightChildIndex, mid + 1, end);
        }
 
        this.treeBuild[stIndex].merge(this.treeBuild[leftChildIndex],
        this.treeBuild[rightChildIndex]);
    }
}
 
// Example usage
const N = 7;
const st = new SegmentTree(N);
 
st.update(2, 5);
console.log(st.query(1, 6));
 
st.update(1, 3);
console.log(st.query(1, 6));


Output

4
3



Last Updated : 16 Nov, 2023
Like Article
Save Article
Previous
Next
Share your thoughts in the comments
Similar Reads