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 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