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

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