2569. Handling Sum Queries After Update


Problem Description

In this problem, you're given two 0-indexed integer arrays, nums1 and nums2, along with a 2D array queries containing various queries that perform operations on these arrays. Each query can be one of three types:

  1. Flip query (Type 1): The query specifies a range [l, r] within nums1. You need to flip the values in this range, meaning that any 0 becomes a 1 and any 1 becomes a 0.

  2. Multiply and Sum query (Type 2): The query specifies a single value p. For each element i in nums2, the value of nums2[i] is incremented by nums1[i] * p.

  3. Sum query (Type 3): This query doesn't require any arguments. It asks for the sum of elements in nums2.

The goal is to process these queries in the order they are given and return an array of the sum values obtained from each third type of query (sum query).

Intuition

To solve this problem, we can observe that type 1 queries can be efficiently processed using a segment tree, which allows us to perform range updates and queries efficiently. A segment tree is a binary tree that is perfect for problems that require us to apply operations over an array segment or find information about an array segment quickly.

For Type 1 queries, we utilize a segment tree to flip values in a range by changing the number of 1s in a range, and we need to handle this lazily (only update when necessary) to ensure efficiency. The lazy propagation technique allows us to avoid unnecessary updates, which is crucial because not every part of the segment tree will be updated with each query.

When handling Type 2 queries, we increase the sum of elements in nums2 by querying the current number of 1s in the range of nums1 using our segment tree and multiplying this by p. We do this instead of iterating over the entire nums1 array, which is inefficient.

Finally, for Type 3 queries, since we are interested in the sum of elements in nums2, we keep a running total of nums2 that is updated with each Type 2 query. The use of the running total allows us to answer Type 3 queries in constant time, as we already have the sum available.

By combining the segment tree for efficient range updates and queries with a running total for nums2, we can process all types of queries efficiently.

Learn more about Segment Tree patterns.

Solution Approach

The given solution uses the concept of a Segment Tree to handle the range updates and queries efficiently. A segment tree is a powerful data structure that allows us to perform operations on array segments (such as summing elements or flipping values) within logarithmic time complexity. Here's how each part of the solution approach is implemented:

  1. Segment Tree Construction:

    • A segment tree is constructed to store the sum of 1s over segments of nums1. Each node of this tree represents a segment of the array, with leaf nodes representing individual elements.
    • The build function recursively constructs the tree. It initializes each node with the number of 1s in the corresponding segment of nums1.
  2. Type 1 Query Handling (Flip Values):

    • To flip values in a range [l, r], the modify function is used. It applies the flip operation lazily to nodes overlapping with [l, r]. When flipping a node representing a segment, the node's sum (s) is updated to be the segment length minus the current sum of 1s (since 0s have become 1s and vice versa).
    • The lazy value is used to delay the update to its children until necessary (when the query directly affects them), which is handled by the pushdown function.
  3. Type 2 Query Handling (Multiply and Sum):

    • For multiplying the sum of 1s in nums1 with a value p and adding it to the sum of nums2, there's no need to iterate through nums1. Instead, a query function obtains the number of 1s in nums1 using the segment tree, which is then multiplied by p and added to the running sum s.
  4. Type 3 Query Handling (Sum Query):

    • Since the running total sum of nums2 is maintained after each Type 2 query, handling Type 3 queries is simple. Each Type 3 query appends the current sum s to the answer list ans, which is what the function eventually returns.
  5. Utility Functions (pushup and pushdown):

    • The pushup function updates a node's sum based on its children's sums, ensuring the correct number of 1s is represented after a range has been modified.
    • The pushdown function is used during the lazy propagation to apply the pending updates to a node's children before making further queries or updates involving those children.

By structuring the solution this way, the problem of processing queries is decomposed into several smaller, more manageable steps. The segment tree handles flipping operations and sum of 1s queries in O(log n) time, where n is the size of nums1, and the Type 2 and Type 3 queries are handled in O(1) time thanks to the running sum. This leads to a very efficient overall solution.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

Which algorithm should you use to find a node that is close to the root of the tree?

Example Walkthrough

Let's demonstrate the solution approach with a small example using the given algorithm.

Consider the following initial data:

1nums1 = [0, 1, 0, 1, 0]
2nums2 = [1, 2, 3, 4, 5]
3queries = [[1, 1, 3], [2, 2], [3], [1, 0, 2], [2, 1], [3]]

Here are the steps we will follow to process these queries:

  1. The segment tree is built from nums1: Each node contains the count of 1s in the corresponding segment. At the start, the tree's representation could look like this:

    1       [1]
    2      /   \
    3    [1]   [1]
    4   / \    / \
    5 [0] [1][0] [1]
    6   |   |  |   |
    7  n1  n2 n3  n4  

    Where [] indicates a node and ni means nums1[i].

  2. Process the first query [1, 1, 3] (Type 1 – Flip query):

    • We need to flip values from index 1 to 3 in nums1. After flipping, nums1 looks like this: [0, 0, 1, 0, 0].
    • The corresponding changes in the segment tree would adjust the counts of 1s in the range [1, 3].
  3. Process the second query [2, 2] (Type 2 – Multiply and Sum query):

    • We read the number of 1s in the range of nums1 using the segment tree (here it's 1) and multiply it by the value p = 2. So the increment for each nums2 element is 1 * 2 = 2.
    • nums2 is updated as follows: [1, 2+2, 3+2, 4+2, 5] = [1, 4, 5, 6, 5].
  4. Process the third query [3] (Type 3 – Sum query):

    • We return the sum of nums2 as [1+4+5+6+5] = 21. The answer starts with [21].
  5. Process the fourth query [1, 0, 2] (Type 1 – Flip query):

    • We flip values from index 0 to 2 in nums1: After flipping, nums1 looks like this: [1, 1, 0, 0, 0].
    • The segment tree is updated to reflect the new counts of 1s accordingly.
  6. Process the fifth query [2, 1] (Type 2 – Multiply and Sum query):

    • We multiply the number of 1s in nums1 (which is 2) by p = 1 and add it to each element in nums2: [1+2, 4+2, 5, 6, 5] = [3, 6, 5, 6, 5].
  7. Process the sixth query [3] (Type 3 – Sum query):

    • The sum of nums2 is calculated again, now [3+6+5+6+5] = 25. The answer list is updated to [21, 25].

By following this approach, the final answer array, corresponding to the sum values obtained from each sum query (Type 3), is [21, 25]. Each Type 1 and Type 2 query modifies the data appropriately, keeping track of the number of 1s in nums1 using a segment tree and maintaining the sum of nums2 for quick access during sum queries.

Solution Implementation

1from typing import List
2
3class Node:
4    def __init__(self):
5        self.left = self.right = 0  # Endpoints of the segment
6        self.sum = self.lazy = 0    # Sum of the segment and lazy mark for lazy propagation
7
8class SegmentTree:
9    def __init__(self, nums: List[int]):
10        self.nums = nums  # Initial array of numbers
11        self.tree = [Node() for _ in range(len(nums) << 2)]  # Segment Tree represented as a list of nodes
12        self.build(1, 1, len(nums))  # Build the Segment Tree
13
14    def build(self, index: int, left: int, right: int):
15        # Recursive function to build the Segment Tree from the initial array
16        self.tree[index].left, self.tree[index].right = left, right
17        if left == right:
18            # Leaf node
19            self.tree[index].sum = self.nums[left - 1]
20            return
21        mid = (left + right) >> 1
22        # Build the left and right subtrees
23        self.build(index << 1, left, mid)
24        self.build(index << 1 | 1, mid + 1, right)
25        # Update the current node's value
26        self.pushup(index)
27
28    def modify(self, index: int, left: int, right: int):
29        # Range update function to modify segments
30        if self.tree[index].left >= left and self.tree[index].right <= right:
31            # If the current segment is entirely within the [left, right] range
32            self.tree[index].lazy ^= 1  # Apply or remove lazy mark
33            # Flip the counts of 1s and 0s in the segment
34            self.tree[index].sum = self.tree[index].right - self.tree[index].left + 1 - self.tree[index].sum
35            return
36        # Push the lazy mark to the children
37        self.pushdown(index)
38        mid = (self.tree[index].left + self.tree[index].right) >> 1
39        # Recursively update the left and/or right subtrees if necessary
40        if left <= mid:
41            self.modify(index << 1, left, right)
42        if right > mid:
43            self.modify(index << 1 | 1, left, right)
44        # After modifications in children, update this node's sum
45        self.pushup(index)
46
47    def query(self, index: int, left: int, right: int):
48        # Range query function to calculate sum in [left, right]
49        if self.tree[index].left >= left and self.tree[index].right <= right:
50            return self.tree[index].sum
51        # Push any lazy marks down the tree before querying children
52        self.pushdown(index)
53        mid = (self.tree[index].left + self.tree[index].right) >> 1
54        res = 0
55        # Recursively query the left and/or right subtrees and accumulate the result
56        if left <= mid:
57            res += self.query(index << 1, left, right)
58        if right > mid:
59            res += self.query(index << 1 | 1, left, right)
60        return res
61
62    def pushup(self, index: int):
63        # Propagates the changes in children up to the current node
64        self.tree[index].sum = self.tree[index << 1].sum + self.tree[index << 1 | 1].sum
65
66    def pushdown(self, index: int):
67        # Pushes the lazy mark from the current node down to its children
68        if self.tree[index].lazy:
69            mid = (self.tree[index].left + self.tree[index].right) >> 1
70            # Apply lazy propagation to the children
71            left_child = index << 1
72            right_child = index << 1 | 1
73            self.tree[left_child].sum = mid - self.tree[index].left + 1 - self.tree[left_child].sum
74            self.tree[left_child].lazy ^= 1
75            self.tree[right_child].sum = self.tree[index].right - mid - self.tree[right_child].sum
76            self.tree[right_child].lazy ^= 1
77          
78            # Remove lazy mark from current node after propagation
79            self.tree[index].lazy ^= 1
80
81class Solution:
82    def handleQuery(self, nums1: List[int], nums2: List[int], queries: List[List[int]]) -> List[int]:
83        tree = SegmentTree(nums1)  # Initialize Segment Tree with nums1
84        sum_of_nums2 = sum(nums2)  # Calculate the sum of nums2
85        answer = []                # List to store answers of each query
86        for operation, a, b in queries:
87            if operation == 1:
88                # Modify operation on the Segment Tree
89                tree.modify(1, a + 1, b + 1)
90            elif operation == 2:
91                # Add to sum_of_nums2 a times the sum of elements in range [1, len(nums1)] in the Segment Tree
92                sum_of_nums2 += a * tree.query(1, 1, len(nums1))
93            else:
94                # Append the current sum_of_nums2 to the answer
95                answer.append(sum_of_nums2)
96        return answer
97
1class Node {
2    int left, right;  // Indicating the range [left, right] this node represents
3    int sum, lazy;    // 'sum' stores the segment sum, 'lazy' is for lazy propagation
4}
5
6class SegmentTree {
7    private Node[] tree;  // The array represents the segment tree
8    private int[] nums;   // The input array based on which segment tree is built
9
10    // Constructor that builds the tree using the provided input array
11    public SegmentTree(int[] nums) {
12        int n = nums.length;
13        this.nums = nums;
14        tree = new Node[n << 2]; // Size of the tree array is 4 times the size of input array
15        for (int i = 0; i < tree.length; ++i) {
16            tree[i] = new Node();
17        }
18        build(1, 1, n); // Start building the tree from the root node
19    }
20
21    // Recursive function to build the segment tree from the given array
22    private void build(int node, int start, int end) {
23        tree[node].left = start;
24        tree[node].right = end;
25        if (start == end) {  // Leaf node condition where the segment is of length 1
26            tree[node].sum = nums[start - 1];
27            return;
28        }
29        int mid = (start + end) >> 1; // Midpoint for splitting the range
30        build(node << 1, start, mid);
31        build(node << 1 | 1, mid + 1, end);
32        pushup(node);
33    }
34
35    // Method to update the segment tree within range [left, right]
36    public void modify(int node, int left, int right) {
37        // If current segment is entirely within range [left, right]
38        if (tree[node].left >= left && tree[node].right <= right) {
39            tree[node].lazy ^= 1; // Toggle the lazy value using XOR operator
40            tree[node].sum = tree[node].right - tree[node].left + 1 - tree[node].sum;
41            return;
42        }
43        pushdown(node); // Process pending updates using lazy propagation
44        int mid = (tree[node].left + tree[node].right) >> 1;
45        // Modify the left child if the query range intersects with it
46        if (left <= mid) {
47            modify(node << 1, left, right);
48        }
49        // Modify the right child if the query range intersects with it
50        if (right > mid) {
51            modify(node << 1 | 1, left, right);
52        }
53        pushup(node); // Update the node with the new values from it's children
54    }
55
56    // Query the sum of a given interval [left, right]
57    public int query(int node, int left, int right) {
58        // If the current segment lies completely within the interval
59        if (tree[node].left >= left && tree[node].right <= right) {
60            return tree[node].sum;
61        }
62        pushdown(node);
63        int mid = (tree[node].left + tree[node].right) >> 1;
64        int result = 0;
65        if (left <= mid) {
66            result += query(node << 1, left, right);
67        }
68        if (right > mid) {
69            result += query(node << 1 | 1, left, right);
70        }
71        return result;
72    }
73
74    // Helper function to update the parent node's sum based on its children's sums
75    private void pushup(int node) {
76        tree[node].sum = tree[node << 1].sum + tree[node << 1 | 1].sum;
77    }
78
79    // Helper function to propagate the lazy update down to the children of a node
80    private void pushdown(int node) {
81        if (tree[node].lazy != 0) {
82            int mid = (tree[node].left + tree[node].right) >> 1;
83            // Apply lazy updates to the left child
84            tree[node << 1].sum = mid - tree[node].left + 1 - tree[node << 1].sum;
85            tree[node << 1].lazy ^= 1;
86            // Apply lazy updates to the right child
87            tree[node << 1 | 1].sum = tree[node].right - mid - tree[node << 1 | 1].sum;
88            tree[node << 1 | 1].lazy ^= 1;
89            // Reset the lazy value in the parent node
90            tree[node].lazy ^= 1;
91        }
92    }
93}
94
95class Solution {
96    // Process queries and return the result as a long array
97    public long[] handleQuery(int[] nums1, int[] nums2, int[][] queries) {
98        SegmentTree tree = new SegmentTree(nums1); // Initialize the segment tree with nums1
99        long totalSum = 0;  // Stores the total sum of elements in nums2
100        for (int x : nums2) {
101            totalSum += x;
102        }
103        int resultCount = 0; // Count of query results to be returned
104        for (int[] q : queries) {
105            if (q[0] == 3) {
106                ++resultCount;
107            }
108        }
109        long[] ans = new long[resultCount]; // Results array
110        int idx = 0; // Index for storing results in 'ans'
111        for (int[] q : queries) { // Process each query in 'queries'
112            if (q[0] == 1) { // Range update query
113                tree.modify(1, q[1] + 1, q[2] + 1);
114            } else if (q[0] == 2) { // Add value to totalSum
115                totalSum += (long) q[1] * tree.query(1, 1, nums2.length);
116            } else { // Query for totalSum
117                ans[idx++] = totalSum;
118            }
119        }
120        return ans; // Return the final result array
121    }
122}
123
1#include <vector>
2using namespace std;
3
4// Definition of Node structure for segment tree.
5class Node {
6public:
7    int start = 0; // Start index of the node's interval.
8    int end = 0; // End index of the node's interval.
9    int sum = 0; // Sum of elements in the node's interval.
10    int lazy = 0; // Lazy propagation tag.
11};
12
13// Definition of the Segment Tree class.
14class SegmentTree {
15public:
16    // Constructor that initializes the segment tree with given values.
17    SegmentTree(vector<int>& nums) {
18        this->nums = nums;
19        int n = nums.size();
20        tree.resize(n << 2); // 4n size to store all nodes.
21        for (int i = 0; i < tree.size(); ++i) {
22            tree[i] = new Node();
23        }
24        build(1, 1, n);
25    }
26
27    // Update function to modify an interval [l, r] with lazy propagation.
28    void modify(int index, int l, int r) {
29        if (tree[index]->start >= l && tree[index]->end <= r) {
30            tree[index]->lazy ^= 1;
31            tree[index]->sum = tree[index]->end - tree[index]->start + 1 - tree[index]->sum;
32            return;
33        }
34        pushdown(index);
35        int mid = (tree[index]->start + tree[index]->end) >> 1;
36        if (l <= mid) {
37            modify(index << 1, l, r);
38        }
39        if (r > mid) {
40            modify(index << 1 | 1, l, r);
41        }
42        pushup(index);
43    }
44
45    // Query function to calculate the sum of the interval [l, r].
46    int query(int index, int l, int r) {
47        if (tree[index]->start >= l && tree[index]->end <= r) {
48            return tree[index]->sum;
49        }
50        pushdown(index);
51        int mid = (tree[index]->start + tree[index]->end) >> 1;
52        int res = 0;
53        if (l <= mid) {
54            res += query(index << 1, l, r);
55        }
56        if (r > mid) {
57            res += query(index << 1 | 1, l, r);
58        }
59        return res;
60    }
61
62private:
63    vector<Node*> tree; // Array of pointers to Node, representing the segment tree.
64    vector<int> nums; // Array of original values used in segment tree.
65
66    // Build function to construct the segment tree from the original array.
67    void build(int index, int start, int end) {
68        tree[index]->start = start;
69        tree[index]->end = end;
70        if (start == end) {
71            tree[index]->sum = nums[start - 1];
72            return;
73        }
74        int mid = (start + end) >> 1;
75        build(index << 1, start, mid);
76        build(index << 1 | 1, mid + 1, end);
77        pushup(index);
78    }
79
80    // Push-up function to update parent node based on the values of child nodes.
81    void pushup(int index) {
82        tree[index]->sum = tree[index << 1]->sum + tree[index << 1 | 1]->sum;
83    }
84
85    // Push-down function to apply the lazy tags to child nodes.
86    void pushdown(int index) {
87        if (tree[index]->lazy) {
88            int mid = (tree[index]->start + tree[index]->end) >> 1;
89            tree[index << 1]->sum = mid - tree[index]->start + 1 - tree[index << 1]->sum;
90            tree[index << 1]->lazy ^= 1;
91            tree[index << 1 | 1]->sum = tree[index]->end - mid - tree[index << 1 | 1]->sum;
92            tree[index << 1 | 1]->lazy ^= 1;
93            tree[index]->lazy = 0; // Reset lazy tag after propagation.
94        }
95    }
96};
97
98// Solution class that uses the segment tree to handle queries.
99class Solution {
100public:
101    // Function to handle queries and return the result as a vector.
102    vector<long long> handleQuery(vector<int>& nums1, vector<int>& nums2, vector<vector<int>>& queries) {
103        SegmentTree* tree = new SegmentTree(nums1); // Initialize the segment tree with nums1.
104        long long totalSum = 0; // Variable to keep track of the total sum during the process.
105        for (int& x : nums2) {
106            totalSum += x;
107        }
108        vector<long long> ans; // Vector to store the answers to the queries.
109        for (auto& query : queries) {
110            if (query[0] == 1) { // If the first element is 1, perform the modify operation.
111                tree->modify(1, query[1] + 1, query[2] + 1);
112            } else if (query[0] == 2) { // If the first element is 2, update the total sum.
113                totalSum += 1LL * query[1] * tree->query(1, 1, nums1.size());
114            } else { // Any other number, add the current total sum to the answer.
115                ans.push_back(totalSum);
116            }
117        }
118        return ans;
119    }
120};
121
1// Using global scope instead of a class structure, define Node as type for segment tree.
2type Node = {
3    start: number; // Starting index of the node's interval.
4    end: number; // Ending index of the node's interval.
5    sum: number; // Sum of the elements in the node's interval.
6    lazy: number; // Lazy propagation tag for deferred updates.
7};
8
9// Initialize an array to hold nodes of the segment tree.
10let tree: Node[];
11
12// Initialize an array to hold original numbers.
13let nums: number[];
14
15// Function to build the segment tree from an array of numbers.
16function build(index: number, start: number, end: number) {
17    tree[index] = {
18        start: start,
19        end: end,
20        sum: 0,
21        lazy: 0
22    };
23
24    if (start === end) {
25        tree[index].sum = nums[start - 1];
26        return;
27    }
28
29    const mid = Math.floor((start + end) / 2);
30    build(index * 2, start, mid);
31    build(index * 2 + 1, mid + 1, end);
32    pushUp(index);
33}
34
35// Function to apply updates from parent nodes to child nodes using lazy propagation.
36function pushDown(index: number) {
37    if (tree[index].lazy) {
38        const mid = Math.floor((tree[index].start + tree[index].end) / 2);
39        const leftChildIndex = index * 2;
40        const rightChildIndex = index * 2 + 1;
41
42        // Apply updates to the left child node.
43        tree[leftChildIndex].sum = mid - tree[index].start + 1 - tree[leftChildIndex].sum;
44        tree[leftChildIndex].lazy ^= 1;
45
46        // Apply updates to the right child node.
47        tree[rightChildIndex].sum = tree[index].end - mid - tree[rightChildIndex].sum;
48        tree[rightChildIndex].lazy ^= 1;
49
50        // Reset lazy propagation tag for the parent.
51        tree[index].lazy = 0;
52    }
53}
54
55// Function to update parent node sums based on children's sums.
56function pushUp(index: number) {
57    const leftChildIndex = index * 2;
58    const rightChildIndex = index * 2 + 1;
59    tree[index].sum = tree[leftChildIndex].sum + tree[rightChildIndex].sum;
60}
61
62// Function to handle updates (toggle) to an interval [l, r] with lazy propagation.
63function modify(index: number, l: number, r: number) {
64    if (tree[index].start >= l && tree[index].end <= r) {
65        tree[index].lazy ^= 1;
66        tree[index].sum = tree[index].end - tree[index].start + 1 - tree[index].sum;
67        return;
68    }
69
70    pushDown(index);
71
72    const mid = Math.floor((tree[index].start + tree[index].end) / 2);
73    if (l <= mid) {
74        modify(index * 2, l, r);
75    }
76
77    if (r > mid) {
78        modify(index * 2 + 1, l, r);
79    }
80
81    pushUp(index);
82}
83
84// Function to query the sum of an interval [l, r].
85function query(index: number, l: number, r: number): number {
86    if (tree[index].start >= l && tree[index].end <= r) {
87        return tree[index].sum;
88    }
89
90    pushDown(index);
91
92    const mid = Math.floor((tree[index].start + tree[index].end) / 2);
93    let res = 0;
94
95    if (l <= mid) {
96        res += query(index * 2, l, r);
97    }
98
99    if (r > mid) {
100        res += query(index * 2 + 1, l, r);
101    }
102
103    return res;
104}
105
106// Function to handle queries on the segment tree and calculate results.
107function handleQuery(nums1: number[], nums2: number[], queries: number[][]): Array<number> {
108    nums = nums1;
109    // Segment tree can have up to 4n nodes for a given array of size n.
110    tree = Array(nums.length * 4); 
111    build(1, 1, nums.length);
112
113    // Keeping track of the overall sum.
114    let totalSum: number = nums2.reduce((acc, val) => acc + val, 0);
115    const answers: Array<number> = [];
116
117    queries.forEach(query => {
118        const [type, param1, param2] = query;
119
120        if (type === 1) {
121            // Perform update/modify operation.
122            modify(1, param1 + 1, param2 + 1);
123        } else if (type === 2) {
124            // Update overall sum using query result.
125            totalSum += query(1, 1, nums.length) * param1;
126        } else {
127            // Push current total sum to the answers for retrieval.