Handling Sum Queries After Update

Hard

#2349Time: O(N + Q * N), where N is the length of the arrays and Q is the number of queries. The initial sum calculations take O(N). Each type 1 query can take up to O(N) time. This makes the overall complexity dominated by the product of the number of queries and the array size.Space: O(A), where A is the number of type 3 queries. This space is used to store the answers. The auxiliary space complexity is O(1).1 company

Data structures

Array Segment Tree

Companies

Trilogy

Prompt

You are given two 0-indexed arrays nums1 and nums2 and a 2D array queries of queries. There are three types of queries:

For a query of type 1, queries[i] = [1, l, r]. Flip the values from 0 to 1 and from 1 to 0 in nums1 from index l to index r. Both l and r are 0-indexed.
For a query of type 2, queries[i] = [2, p, 0]. For every index 0 <= i < n, set nums2[i] = nums2[i] + nums1[i] * p.
For a query of type 3, queries[i] = [3, 0, 0]. Find the sum of the elements in nums2.

Return an array containing all the answers to the third type queries.

Example 1:

Input: nums1 = [1,0,1], nums2 = [0,0,0], queries = [[1,1,1],[2,1,0],[3,0,0]]
Output: [3]
Explanation: After the first query nums1 becomes [1,1,1]. After the second query, nums2 becomes [1,1,1], so the answer to the third query is 3. Thus, [3] is returned.

Example 2:

Input: nums1 = [1], nums2 = [5], queries = [[2,0,0],[3,0,0]]
Output: [5]
Explanation: After the first query, nums2 remains [5], so the answer to the second query is 5. Thus, [5] is returned.

Constraints:

1 <= nums1.length,nums2.length <= 10⁵
nums1.length = nums2.length
1 <= queries.length <= 10⁵
queries[i].length = 3
0 <= l <= r <= nums1.length - 1
0 <= p <= 10⁶
0 <= nums1[i] <= 1
0 <= nums2[i] <= 10⁹

Approaches

2 approaches with complexity analysis and trade-offs.

This approach directly simulates the process described in the problem. We maintain the arrays nums1 and nums2 and process each query one by one. To make it slightly more efficient than a pure naive simulation, we keep track of the total sum of nums2 and the count of 1s in nums1 in separate variables, updating them as queries are processed. This avoids re-calculating sums from scratch every time.

Algorithm

Initialize a long variable sum2 with the total sum of nums2.
Initialize an int or long variable sum1 with the count of 1s in nums1.
Create a list to store the answers for type 3 queries.
Iterate through each query:
- For a type 1 query [1, l, r], iterate from l to r in nums1. Flip each element and update sum1 accordingly.
- For a type 2 query [2, p, 0], update sum2 by adding p * sum1.
- For a type 3 query [3, 0, 0], add the current sum2 to the answer list.
Convert the answer list to an array and return it.

Walkthrough

In this method, we first pre-calculate the initial sum of nums2 and the initial count of 1s in nums1. We store these in sum2 (as a long to prevent overflow) and sum1 respectively. Then, we process the queries sequentially.

Type 1 Query [1, l, r]: We loop through the specified range [l, r] of nums1. For each element, we check its current value. If it's 1, we decrement sum1; if it's 0, we increment sum1. Then we flip the value (nums1[i] = 1 - nums1[i]). This operation has a time complexity proportional to the length of the range, which is O(N) in the worst case.
Type 2 Query [2, p, 0]: We use the pre-computed sum1. The total sum of nums2 increases by p for each 1 in nums1. So, we update sum2 by adding (long)p * sum1. This is an O(1) operation.
Type 3 Query [3, 0, 0]: We simply take the current value of sum2 and add it to our list of answers. This is also an O(1) operation.

After processing all queries, we return the collected answers.

1class Solution {2    public long[] handleQuery(int[] nums1, int[] nums2, int[][] queries) {3        int n = nums1.length;4        long sum1 = 0;5        for (int x : nums1) {6            sum1 += x;7        }8 9        long sum2 = 0;10        for (int x : nums2) {11            sum2 += x;12        }13 14        java.util.List<Long> answers = new java.util.ArrayList<>();15 16        for (int[] query : queries) {17            int type = query[0];18            if (type == 1) {19                int l = query[1];20                int r = query[2];21                for (int i = l; i <= r; i++) {22                    if (nums1[i] == 1) {23                        sum1--;24                        nums1[i] = 0;25                    } else {26                        sum1++;27                        nums1[i] = 1;28                    }29                }30            } else if (type == 2) {31                long p = query[1];32                sum2 += p * sum1;33            } else { // type == 334                answers.add(sum2);35            }36        }37 38        long[] result = new long[answers.size()];39        for (int i = 0; i < answers.size(); i++) {40            result[i] = answers.get(i);41        }42        return result;43    }44}

Complexity

Time

O(N + Q * N), where N is the length of the arrays and Q is the number of queries. The initial sum calculations take O(N). Each type 1 query can take up to O(N) time. This makes the overall complexity dominated by the product of the number of queries and the array size.

Space

O(A), where A is the number of type 3 queries. This space is used to store the answers. The auxiliary space complexity is O(1).

Trade-offs

Pros

Simple to understand and straightforward to implement.
Low memory overhead, using O(1) extra space (excluding the answer list).

Cons

Highly inefficient for large inputs due to the O(N) complexity of type 1 queries.
Will likely result in a Time Limit Exceeded (TLE) error on platforms with strict time limits.

The bottleneck in the brute-force approach is the type 1 query, which involves a range update on nums1. This type of operation is a classic use case for a Segment Tree data structure with lazy propagation. By using a Segment Tree to maintain the count of 1s in nums1, we can reduce the time complexity of range flips from O(N) to O(log N), making the overall solution efficient enough to pass within the given constraints.

Algorithm

Initialize sum2 with the sum of nums2.
Build a Segment Tree over nums1. Each node stores the count of 1s in its range. The tree should also support lazy propagation for flip operations.
Create a list to store answers.
Loop through each query:
- If query type is 1 ([1, l, r]): Perform a range update on the Segment Tree for [l, r]. This will take O(log N) time.
- If query type is 2 ([2, p, 0]): Get the total count of 1s (sum1) from the root of the Segment Tree. Update sum2 += (long)p * sum1.
- If query type is 3 ([3, 0, 0]): Add sum2 to the answer list.
Return the answer list.

Walkthrough

This approach optimizes the handling of type 1 queries using a Segment Tree.

Data Structure: We build a Segment Tree on nums1. Each node in the tree stores the count of 1s within its corresponding array segment. We also use a boolean array, lazy, of the same size as the tree to handle range flips efficiently. A lazy[node] = true flag indicates that the segment corresponding to node needs to be flipped.

Algorithm Steps:

Initialization: Calculate the initial sum of nums2 and store it in a long variable sum2. Build the Segment Tree from nums1 in O(N) time. The root of the tree will hold the total initial count of 1s.
Type 1 Query [1, l, r]: This is a range flip. We call an updateRange function on our Segment Tree. This function traverses the tree and, for segments fully contained in [l, r], it updates the count (new_count = length - old_count) and toggles the lazy flag for its children. This operation takes O(log N) time.
Type 2 Query [2, p, 0]: We need the total count of 1s in nums1. This value is stored at the root of our Segment Tree (tree[1]). We retrieve this value in O(1) time and update sum2 by adding (long)p * tree[1].
Type 3 Query [3, 0, 0]: We add the current value of sum2 to our list of answers in O(1) time.

This brings the total time complexity down significantly, making it suitable for the problem's constraints.

1class Solution {2    long[] tree;3    boolean[] lazy;4    int n;5 6    private void build(int node, int start, int end, int[] nums1) {7        if (start == end) {8            tree[node] = nums1[start];9            return;10        }11        int mid = start + (end - start) / 2;12        build(2 * node, start, mid, nums1);13        build(2 * node + 1, mid + 1, end, nums1);14        tree[node] = tree[2 * node] + tree[2 * node + 1];15    }16 17    private void push(int node, int start, int end) {18        if (lazy[node]) {19            long rangeLen = end - start + 1;20            tree[node] = rangeLen - tree[node];21            if (start != end) {22                lazy[2 * node] = !lazy[2 * node];23                lazy[2 * node + 1] = !lazy[2 * node + 1];24            }25            lazy[node] = false;26        }27    }28 29    private void updateRange(int node, int start, int end, int l, int r) {30        push(node, start, end);31        if (start > end || start > r || end < l) {32            return;33        }34        if (l <= start && end <= r) {35            lazy[node] = !lazy[node];36            push(node, start, end);37            return;38        }39        int mid = start + (end - start) / 2;40        updateRange(2 * node, start, mid, l, r);41        updateRange(2 * node + 1, mid + 1, end, l, r);42        43        push(2*node, start, mid);44        push(2*node+1, mid+1, end);45        tree[node] = tree[2 * node] + tree[2 * node + 1];46    }47 48    public long[] handleQuery(int[] nums1, int[] nums2, int[][] queries) {49        this.n = nums1.length;50        this.tree = new long[4 * n];51        this.lazy = new boolean[4 * n];52 53        build(1, 0, n - 1, nums1);54 55        long sum2 = 0;56        for (int num : nums2) {57            sum2 += num;58        }59 60        java.util.List<Long> ansList = new java.util.ArrayList<>();61        for (int[] query : queries) {62            if (query[0] == 1) {63                updateRange(1, 0, n - 1, query[1], query[2]);64            } else if (query[0] == 2) {65                long p = query[1];66                push(1, 0, n - 1); // Ensure root is up-to-date67                sum2 += p * tree[1];68            } else { // query[0] == 369                ansList.add(sum2);70            }71        }72 73        long[] result = new long[ansList.size()];74        for (int i = 0; i < ansList.size(); i++) {75            result[i] = ansList.get(i);76        }77        return result;78    }79}

Complexity

Time

O(N + Q * log N), where N is the length of the arrays and Q is the number of queries. Building the tree takes O(N). Each type 1 query takes O(log N), while type 2 and 3 queries take O(1). This is a significant improvement over the brute-force approach.

Space

O(N + A), where N is the array size and A is the number of type 3 queries. The Segment Tree and lazy array require O(N) space, and the answer list requires O(A) space.

Trade-offs

Pros

Very efficient, with a logarithmic time complexity for the most expensive operation.
Can handle large inputs and pass within typical time limits.
It's a standard and powerful technique applicable to many similar range-based problems.

Cons

More complex to implement correctly compared to the brute-force approach.
The implementation of lazy propagation can be tricky and prone to bugs.

Solutions

Solution

1class Node { int l , r ; int s , lazy ; } class SegmentTree { private Node [] tr ; private int [] nums ; public SegmentTree ( int [] nums ) { int n = nums . length ; this . nums = nums ; tr = new Node [ n << 2 ]; for ( int i = 0 ; i < tr . length ; ++ i ) { tr [ i ] = new Node (); } build ( 1 , 1 , n ); } private void build ( int u , int l , int r ) { tr [ u ]. l = l ; tr [ u ]. r = r ; if ( l == r ) { tr [ u ]. s = nums [ l - 1 ]; return ; } int mid = ( l + r ) >> 1 ; build ( u << 1 , l , mid ); build ( u << 1 | 1 , mid + 1 , r ); pushup ( u ); } public void modify ( int u , int l , int r ) { if ( tr [ u ]. l >= l && tr [ u ]. r <= r ) { tr [ u ]. lazy ^= 1 ; tr [ u ]. s = tr [ u ]. r - tr [ u ]. l + 1 - tr [ u ]. s ; return ; } pushdown ( u ); int mid = ( tr [ u ]. l + tr [ u ]. r ) >> 1 ; if ( l <= mid ) { modify ( u << 1 , l , r ); } if ( r > mid ) { modify ( u << 1 | 1 , l , r ); } pushup ( u ); } public int query ( int u , int l , int r ) { if ( tr [ u ]. l >= l && tr [ u ]. r <= r ) { return tr [ u ]. s ; } pushdown ( u ); int mid = ( tr [ u ]. l + tr [ u ]. r ) >> 1 ; int res = 0 ; if ( l <= mid ) { res += query ( u << 1 , l , r ); } if ( r > mid ) { res += query ( u << 1 | 1 , l , r ); } return res ; } private void pushup ( int u ) { tr [ u ]. s = tr [ u << 1 ]. s + tr [ u << 1 | 1 ]. s ; } private void pushdown ( int u ) { if ( tr [ u ]. lazy == 1 ) { int mid = ( tr [ u ]. l + tr [ u ]. r ) >> 1 ; tr [ u << 1 ]. s = mid - tr [ u ]. l + 1 - tr [ u << 1 ]. s ; tr [ u << 1 ]. lazy ^= 1 ; tr [ u << 1 | 1 ]. s = tr [ u ]. r - mid - tr [ u << 1 | 1 ]. s ; tr [ u << 1 | 1 ]. lazy ^= 1 ; tr [ u ]. lazy ^= 1 ; } } } class Solution { public long [] handleQuery ( int [] nums1 , int [] nums2 , int [][] queries ) { SegmentTree tree = new SegmentTree ( nums1 ); long s = 0 ; for ( int x : nums2 ) { s += x ; } int m = 0 ; for ( var q : queries ) { if ( q [ 0 ] == 3 ) { ++ m ; } } long [] ans = new long [ m ]; int i = 0 ; for ( var q : queries ) { if ( q [ 0 ] == 1 ) { tree . modify ( 1 , q [ 1 ] + 1 , q [ 2 ] + 1 ); } else if ( q [ 0 ] == 2 ) { s += 1L * q [ 1 ] * tree . query ( 1 , 1 , nums2 . length ); } else { ans [ i ++] = s ; } } return ans ; } }

1class Node { public: int l = 0 , r = 0 ; int s = 0 , lazy = 0 ; }; class SegmentTree { public: SegmentTree ( vector < int >& nums ) { this -> nums = nums ; int n = nums . size (); tr . resize ( n << 2 ); for ( int i = 0 ; i < tr . size (); ++ i ) { tr [ i ] = new Node (); } build ( 1 , 1 , n ); } void modify ( int u , int l , int r ) { if ( tr [ u ] -> l >= l && tr [ u ] -> r <= r ) { tr [ u ] -> lazy ^= 1 ; tr [ u ] -> s = tr [ u ] -> r - tr [ u ] -> l + 1 - tr [ u ] -> s ; return ; } pushdown ( u ); int mid = ( tr [ u ] -> l + tr [ u ] -> r ) >> 1 ; if ( l <= mid ) { modify ( u << 1 , l , r ); } if ( r > mid ) { modify ( u << 1 | 1 , l , r ); } pushup ( u ); } int query ( int u , int l , int r ) { if ( tr [ u ] -> l >= l && tr [ u ] -> r <= r ) { return tr [ u ] -> s ; } pushdown ( u ); int mid = ( tr [ u ] -> l + tr [ u ] -> r ) >> 1 ; int res = 0 ; if ( l <= mid ) { res += query ( u << 1 , l , r ); } if ( r > mid ) { res += query ( u << 1 | 1 , l , r ); } return res ; } private: vector < Node *> tr ; vector < int > nums ; void build ( int u , int l , int r ) { tr [ u ] -> l = l ; tr [ u ] -> r = r ; if ( l == r ) { tr [ u ] -> s = nums [ l - 1 ]; return ; } int mid = ( l + r ) >> 1 ; build ( u << 1 , l , mid ); build ( u << 1 | 1 , mid + 1 , r ); pushup ( u ); } void pushup ( int u ) { tr [ u ] -> s = tr [ u << 1 ] -> s + tr [ u << 1 | 1 ] -> s ; } void pushdown ( int u ) { if ( tr [ u ] -> lazy ) { int mid = ( tr [ u ] -> l + tr [ u ] -> r ) >> 1 ; tr [ u << 1 ] -> s = mid - tr [ u ] -> l + 1 - tr [ u << 1 ] -> s ; tr [ u << 1 ] -> lazy ^= 1 ; tr [ u << 1 | 1 ] -> s = tr [ u ] -> r - mid - tr [ u << 1 | 1 ] -> s ; tr [ u << 1 | 1 ] -> lazy ^= 1 ; tr [ u ] -> lazy ^= 1 ; } } }; class Solution { public: vector < long long > handleQuery ( vector < int >& nums1 , vector < int >& nums2 , vector < vector < int >>& queries ) { SegmentTree * tree = new SegmentTree ( nums1 ); long long s = 0 ; for ( int & x : nums2 ) { s += x ; } vector < long long > ans ; for ( auto & q : queries ) { if ( q [ 0 ] == 1 ) { tree -> modify ( 1 , q [ 1 ] + 1 , q [ 2 ] + 1 ); } else if ( q [ 0 ] == 2 ) { s += 1LL * q [ 1 ] * tree -> query ( 1 , 1 , nums1 . size ()); } else { ans . push_back ( s ); } } return ans ; } };

1class Node : def __init__ ( self ): self . l = self . r = 0 self . s = self . lazy = 0 class SegmentTree : def __init__ ( self , nums ): self . nums = nums n = len ( nums ) self . tr = [ Node () for _ in range ( n << 2 )] self . build ( 1 , 1 , n ) def build ( self , u , l , r ): self . tr [ u ]. l , self . tr [ u ]. r = l , r if l == r : self . tr [ u ]. s = self . nums [ l - 1 ] return mid = ( l + r ) >> 1 self . build ( u << 1 , l , mid ) self . build ( u << 1 | 1 , mid + 1 , r ) self . pushup ( u ) def modify ( self , u , l , r ): if self . tr [ u ]. l >= l and self . tr [ u ]. r <= r : self . tr [ u ]. lazy ^= 1 self . tr [ u ]. s = self . tr [ u ]. r - self . tr [ u ]. l + 1 - self . tr [ u ]. s return self . pushdown ( u ) mid = ( self . tr [ u ]. l + self . tr [ u ]. r ) >> 1 if l <= mid : self . modify ( u << 1 , l , r ) if r > mid : self . modify ( u << 1 | 1 , l , r ) self . pushup ( u ) def query ( self , u , l , r ): if self . tr [ u ]. l >= l and self . tr [ u ]. r <= r : return self . tr [ u ]. s self . pushdown ( u ) mid = ( self . tr [ u ]. l + self . tr [ u ]. r ) >> 1 res = 0 if l <= mid : res += self . query ( u << 1 , l , r ) if r > mid : res += self . query ( u << 1 | 1 , l , r ) return res def pushup ( self , u ): self . tr [ u ]. s = self . tr [ u << 1 ]. s + self . tr [ u << 1 | 1 ]. s def pushdown ( self , u ): if self . tr [ u ]. lazy : mid = ( self . tr [ u ]. l + self . tr [ u ]. r ) >> 1 self . tr [ u << 1 ]. s = mid - self . tr [ u ]. l + 1 - self . tr [ u << 1 ]. s self . tr [ u << 1 ]. lazy ^= 1 self . tr [ u << 1 | 1 ]. s = self . tr [ u ]. r - mid - self . tr [ u << 1 | 1 ]. s self . tr [ u << 1 | 1 ]. lazy ^= 1 self . tr [ u ]. lazy ^= 1 class Solution : def handleQuery ( self , nums1 : List [ int ], nums2 : List [ int ], queries : List [ List [ int ]] ) -> List [ int ]: tree = SegmentTree ( nums1 ) s = sum ( nums2 ) ans = [] for op , a , b in queries : if op == 1 : tree . modify ( 1 , a + 1 , b + 1 ) elif op == 2 : s += a * tree . query ( 1 , 1 , len ( nums1 )) else : ans . append ( s ) return ans

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem