Facebook Pixel

2569. Handling Sum Queries After Update

HardSegment TreeArray
Leetcode Link

Problem Description

You are given two arrays nums1 and nums2 (both 0-indexed) and need to process a series of queries. Each query can be one of three types:

Type 1 Query: queries[i] = [1, l, r]

  • Flip all values in nums1 from index l to index r (inclusive)
  • This means changing all 0s to 1s and all 1s to 0s in that range

Type 2 Query: queries[i] = [2, p, 0]

  • Update every element in nums2 using the formula: nums2[i] = nums2[i] + nums1[i] * p
  • This adds p times the corresponding nums1 value to each element of nums2

Type 3 Query: queries[i] = [3, 0, 0]

  • Calculate and return the sum of all elements in nums2

Your task is to process all queries in order and return an array containing the results of all Type 3 queries.

The key insight is that Type 2 operations add p * sum(nums1) to the total sum of nums2, since we're adding p * nums1[i] to each nums2[i]. This means we need to efficiently handle range flips in nums1 (Type 1) and track its sum to properly update the running total of nums2 (Type 2), without explicitly updating every element of nums2 each time.

Quick Interview Experience
Help others by sharing your interview experience
Have you seen this problem before?

Intuition

Let's think about what each operation really does:

For Type 3 queries, we need to return the sum of nums2. Instead of recalculating this sum every time, we can maintain a running total s that represents sum(nums2).

For Type 2 queries, notice that when we perform nums2[i] = nums2[i] + nums1[i] * p for all indices, we're essentially adding p * sum(nums1) to our running total s. This is because:

  • new_sum(nums2) = sum(nums2[i] + nums1[i] * p)
  • = sum(nums2[i]) + sum(nums1[i] * p)
  • = sum(nums2) + p * sum(nums1)

So we don't need to update individual elements of nums2 - we just need to know the sum of nums1 at any given moment.

For Type 1 queries, we flip bits in a range of nums1. When a bit flips from 0 to 1, it contributes +1 to the sum, and when it flips from 1 to 0, it contributes -1. For a range [l, r]:

  • The new sum in that range = (r - l + 1) - old_sum_in_range
  • This is because each 1 becomes 0 and each 0 becomes 1

Since Type 1 involves range updates (flipping) and we need to efficiently query the sum of nums1, this naturally leads us to use a segment tree with lazy propagation. The segment tree allows us to:

  • Flip a range of values in O(log n) time
  • Query the total sum in O(log n) time
  • Use lazy propagation to defer updates until necessary, making range updates efficient

The lazy tag in each node indicates whether the range represented by that node needs to be flipped. When we propagate the lazy tag down, we flip the child nodes and toggle their lazy tags.

Learn more about Segment Tree patterns.

Solution Approach

The solution uses a Segment Tree with Lazy Propagation to efficiently handle range updates and queries on nums1.

Node Structure

Each node in the segment tree contains:

  • l, r: The left and right endpoints of the range (1-indexed)
  • s: The sum of elements in this range
  • lazy: A flag indicating if this range needs to be flipped

Building the Segment Tree

The build(u, l, r) function constructs the tree recursively:

  • If l == r, it's a leaf node representing a single element from nums1
  • Otherwise, split the range at midpoint mid = (l + r) >> 1
  • Build left child at index u << 1 for range [l, mid]
  • Build right child at index u << 1 | 1 for range [mid + 1, r]
  • Update parent's sum using pushup(u)

Range Modification (Flipping)

The modify(u, l, r) function flips values in range [l, r]:

  • If current node's range is completely within [l, r]:
    • Toggle the lazy flag: lazy ^= 1
    • Update the sum: s = (r - l + 1) - s (since all 0s become 1s and vice versa)
  • Otherwise, push down lazy tag and recursively modify children

Lazy Propagation

The pushdown(u) function propagates lazy tags to children:

  • If current node has a lazy tag:
    • Flip left child's sum: left.s = (mid - l + 1) - left.s
    • Toggle left child's lazy: left.lazy ^= 1
    • Do the same for right child
    • Clear current node's lazy tag

Range Query

The query(u, l, r) function returns the sum of range [l, r]:

  • If current node's range is completely within [l, r], return its sum
  • Otherwise, push down lazy tag and recursively query children
  • Combine results from left and right children

Main Algorithm

  1. Initialize segment tree with nums1
  2. Calculate initial sum s = sum(nums2)
  3. Process each query:
    • Type 1 [1, l, r]: Call modify(1, l+1, r+1) to flip range (converting to 1-indexed)
    • Type 2 [2, p, 0]: Update s += p * query(1, 1, n) where query returns current sum of nums1
    • Type 3 [3, 0, 0]: Append current value of s to answer array

The time complexity is O(q * log n) where q is the number of queries and n is the length of nums1. The space complexity is O(n) for the segment tree.

Ready to land your dream job?

Unlock your dream job with a 5-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's walk through a small example to illustrate how the solution works.

Input:

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

Initial Setup:

  1. Build segment tree for nums1 = [1, 0, 1]:

    • Root node (range [1,3]): sum = 2
    • Left child (range [1,2]): sum = 1
    • Right child (range [3,3]): sum = 1
    • Leaf nodes: [1], [0], [1]

  2. Calculate initial sum: s = sum(nums2) = 2 + 3 + 4 = 9

Processing Queries:

Query 1: [1, 0, 1] - Flip nums1[0:1] (indices 0 and 1)

  • Convert to 1-indexed: flip range [1, 2]
  • Before flip: nums1 = [1, 0, 1], segment tree sum = 2
  • After flip: nums1 = [0, 1, 1], segment tree sum = 2
  • The segment tree handles this by:
    • Node [1,2] gets lazy flag toggled
    • Its sum changes from 1 to (2-1) = 1 (but different elements are 1s now)

Query 2: [2, 2, 0] - Add 2 * nums1[i] to each nums2[i]

  • Current sum of nums1 = query(1, 1, 3) = 2
  • Update running total: s = 9 + 2 * 2 = 13
  • This represents: nums2 = [2+0*2, 3+1*2, 4+1*2] = [2, 5, 6]

Query 3: [3, 0, 0] - Return sum of nums2

  • Return current s = 13
  • Add 13 to answer array: ans = [13]

Query 4: [1, 1, 2] - Flip nums1[1:2] (indices 1 and 2)

  • Convert to 1-indexed: flip range [2, 3]
  • Before flip: nums1 = [0, 1, 1], segment tree sum = 2
  • After flip: nums1 = [0, 0, 0], segment tree sum = 0
  • The segment tree handles this by:
    • Nodes [2,2] and [3,3] get updated
    • Total sum changes from 2 to 0

Query 5: [3, 0, 0] - Return sum of nums2

  • Return current s = 13 (unchanged since last Type 2 query)
  • Add 13 to answer array: ans = [13, 13]

Final Answer: [13, 13]

Key Insights from Example:

  1. We never actually modify the nums2 array - we just track its sum in variable s
  2. Type 1 queries efficiently update ranges in the segment tree without iterating through elements
  3. Type 2 queries only need the total sum of nums1, which the segment tree provides in O(log n) time
  4. The lazy propagation ensures we don't update nodes until necessary, making range flips efficient

Solution Implementation

1from typing import List
2
3
4class Node:
5    """Represents a node in the segment tree."""
6  
7    def __init__(self):
8        self.left = 0          # Left boundary of the segment
9        self.right = 0         # Right boundary of the segment
10        self.sum = 0           # Sum of 1s in the segment
11        self.lazy = 0          # Lazy propagation flag for flip operations
12
13
14class SegmentTree:
15    """
16    Segment tree implementation with lazy propagation for range updates.
17    Supports range flip operations and range sum queries.
18    """
19  
20    def __init__(self, nums: List[int]):
21        """
22        Initialize the segment tree with the given array.
23      
24        Args:
25            nums: Initial array of 0s and 1s
26        """
27        self.nums = nums
28        n = len(nums)
29        # Create array of nodes with size 4*n for the segment tree
30        self.tree = [Node() for _ in range(n << 2)]
31        self.build(1, 1, n)
32  
33    def build(self, node_idx: int, left: int, right: int) -> None:
34        """
35        Build the segment tree recursively.
36      
37        Args:
38            node_idx: Current node index in the tree
39            left: Left boundary of current segment
40            right: Right boundary of current segment
41        """
42        self.tree[node_idx].left = left
43        self.tree[node_idx].right = right
44      
45        # Leaf node: directly assign the value
46        if left == right:
47            self.tree[node_idx].sum = self.nums[left - 1]  # Convert to 0-indexed
48            return
49      
50        # Build left and right subtrees
51        mid = (left + right) >> 1
52        self.build(node_idx << 1, left, mid)
53        self.build(node_idx << 1 | 1, mid + 1, right)
54      
55        # Update current node's sum based on children
56        self.pushup(node_idx)
57  
58    def modify(self, node_idx: int, left: int, right: int) -> None:
59        """
60        Flip all bits in the range [left, right].
61      
62        Args:
63            node_idx: Current node index in the tree
64            left: Left boundary of the range to flip
65            right: Right boundary of the range to flip
66        """
67        # Current segment is completely within the update range
68        if self.tree[node_idx].left >= left and self.tree[node_idx].right <= right:
69            # Toggle lazy flag
70            self.tree[node_idx].lazy ^= 1
71            # Flip all bits: new sum = total elements - old sum
72            segment_size = self.tree[node_idx].right - self.tree[node_idx].left + 1
73            self.tree[node_idx].sum = segment_size - self.tree[node_idx].sum
74            return
75      
76        # Push down lazy updates before proceeding
77        self.pushdown(node_idx)
78      
79        # Recursively update overlapping segments
80        mid = (self.tree[node_idx].left + self.tree[node_idx].right) >> 1
81        if left <= mid:
82            self.modify(node_idx << 1, left, right)
83        if right > mid:
84            self.modify(node_idx << 1 | 1, left, right)
85      
86        # Update current node's sum after modifying children
87        self.pushup(node_idx)
88  
89    def query(self, node_idx: int, left: int, right: int) -> int:
90        """
91        Query the sum of elements in range [left, right].
92      
93        Args:
94            node_idx: Current node index in the tree
95            left: Left boundary of the query range
96            right: Right boundary of the query range
97          
98        Returns:
99            Sum of 1s in the specified range
100        """
101        # Current segment is completely within the query range
102        if self.tree[node_idx].left >= left and self.tree[node_idx].right <= right:
103            return self.tree[node_idx].sum
104      
105        # Push down lazy updates before querying
106        self.pushdown(node_idx)
107      
108        # Query overlapping segments
109        mid = (self.tree[node_idx].left + self.tree[node_idx].right) >> 1
110        result = 0
111        if left <= mid:
112            result += self.query(node_idx << 1, left, right)
113        if right > mid:
114            result += self.query(node_idx << 1 | 1, left, right)
115      
116        return result
117  
118    def pushup(self, node_idx: int) -> None:
119        """
120        Update parent node's sum based on its children.
121      
122        Args:
123            node_idx: Current node index to update
124        """
125        left_child = node_idx << 1
126        right_child = node_idx << 1 | 1
127        self.tree[node_idx].sum = self.tree[left_child].sum + self.tree[right_child].sum
128  
129    def pushdown(self, node_idx: int) -> None:
130        """
131        Push lazy updates to children nodes.
132      
133        Args:
134            node_idx: Current node index with lazy updates
135        """
136        if self.tree[node_idx].lazy:
137            mid = (self.tree[node_idx].left + self.tree[node_idx].right) >> 1
138            left_child = node_idx << 1
139            right_child = node_idx << 1 | 1
140          
141            # Update left child
142            left_size = mid - self.tree[node_idx].left + 1
143            self.tree[left_child].sum = left_size - self.tree[left_child].sum
144            self.tree[left_child].lazy ^= 1
145          
146            # Update right child
147            right_size = self.tree[node_idx].right - mid
148            self.tree[right_child].sum = right_size - self.tree[right_child].sum
149            self.tree[right_child].lazy ^= 1
150          
151            # Clear parent's lazy flag
152            self.tree[node_idx].lazy = 0
153
154
155class Solution:
156    """Solution for handling three types of queries on two arrays."""
157  
158    def handleQuery(
159        self, nums1: List[int], nums2: List[int], queries: List[List[int]]
160    ) -> List[int]:
161        """
162        Process queries on two arrays.
163      
164        Query types:
165        1. Flip bits in nums1 from index a to b
166        2. Add a * sum(nums1) to all elements in nums2
167        3. Return sum of nums2
168      
169        Args:
170            nums1: Binary array (0s and 1s)
171            nums2: Integer array
172            queries: List of [operation, a, b] queries
173          
174        Returns:
175            List of results for type 3 queries
176        """
177        # Initialize segment tree for nums1
178        tree = SegmentTree(nums1)
179      
180        # Track running sum of nums2
181        current_sum = sum(nums2)
182      
183        # Store results for type 3 queries
184        results = []
185      
186        for operation, param_a, param_b in queries:
187            if operation == 1:
188                # Flip bits in range [param_a, param_b]
189                # Convert to 1-indexed for segment tree
190                tree.modify(1, param_a + 1, param_b + 1)
191            elif operation == 2:
192                # Add param_a * sum(nums1) to all elements in nums2
193                # This increases total sum by param_a * sum(nums1) * len(nums2)
194                ones_count = tree.query(1, 1, len(nums1))
195                current_sum += param_a * ones_count
196            else:  # operation == 3
197                # Return current sum of nums2
198                results.append(current_sum)
199      
200        return results
201
1/**
2 * Node class represents a node in the segment tree
3 */
4class Node {
5    int left, right;     // Range boundaries [left, right]
6    int sum;             // Sum of 1s in the range
7    int lazy;            // Lazy propagation flag (0 or 1)
8}
9
10/**
11 * Segment Tree implementation with lazy propagation for range flip operations
12 */
13class SegmentTree {
14    private Node[] tree;
15    private int[] nums;
16  
17    /**
18     * Constructor to initialize the segment tree with given array
19     * @param nums - input array
20     */
21    public SegmentTree(int[] nums) {
22        int n = nums.length;
23        this.nums = nums;
24        // Allocate 4 times the array size for segment tree nodes
25        tree = new Node[n << 2];
26        for (int i = 0; i < tree.length; ++i) {
27            tree[i] = new Node();
28        }
29        // Build the tree starting from root (index 1), covering range [1, n]
30        build(1, 1, n);
31    }
32  
33    /**
34     * Recursively build the segment tree
35     * @param nodeIndex - current node index in tree array
36     * @param left - left boundary of current range
37     * @param right - right boundary of current range
38     */
39    private void build(int nodeIndex, int left, int right) {
40        tree[nodeIndex].left = left;
41        tree[nodeIndex].right = right;
42      
43        // Leaf node: directly assign value from nums array
44        if (left == right) {
45            tree[nodeIndex].sum = nums[left - 1];  // Convert to 0-indexed
46            return;
47        }
48      
49        // Build left and right subtrees
50        int mid = (left + right) >> 1;
51        build(nodeIndex << 1, left, mid);
52        build(nodeIndex << 1 | 1, mid + 1, right);
53      
54        // Update current node's sum based on children
55        pushup(nodeIndex);
56    }
57  
58    /**
59     * Modify (flip) values in range [left, right]
60     * @param nodeIndex - current node index
61     * @param left - left boundary of query range
62     * @param right - right boundary of query range
63     */
64    public void modify(int nodeIndex, int left, int right) {
65        // Current node's range is completely within query range
66        if (tree[nodeIndex].left >= left && tree[nodeIndex].right <= right) {
67            // Toggle lazy flag
68            tree[nodeIndex].lazy ^= 1;
69            // Flip all bits: new sum = range_size - old_sum
70            tree[nodeIndex].sum = tree[nodeIndex].right - tree[nodeIndex].left + 1 - tree[nodeIndex].sum;
71            return;
72        }
73      
74        // Push down lazy updates before processing children
75        pushdown(nodeIndex);
76      
77        int mid = (tree[nodeIndex].left + tree[nodeIndex].right) >> 1;
78        // Process left child if query range overlaps with left half
79        if (left <= mid) {
80            modify(nodeIndex << 1, left, right);
81        }
82        // Process right child if query range overlaps with right half
83        if (right > mid) {
84            modify(nodeIndex << 1 | 1, left, right);
85        }
86      
87        // Update current node based on children
88        pushup(nodeIndex);
89    }
90  
91    /**
92     * Query the sum of values in range [left, right]
93     * @param nodeIndex - current node index
94     * @param left - left boundary of query range
95     * @param right - right boundary of query range
96     * @return sum of values in the range
97     */
98    public int query(int nodeIndex, int left, int right) {
99        // Current node's range is completely within query range
100        if (tree[nodeIndex].left >= left && tree[nodeIndex].right <= right) {
101            return tree[nodeIndex].sum;
102        }
103      
104        // Push down lazy updates before processing children
105        pushdown(nodeIndex);
106      
107        int mid = (tree[nodeIndex].left + tree[nodeIndex].right) >> 1;
108        int result = 0;
109      
110        // Query left child if needed
111        if (left <= mid) {
112            result += query(nodeIndex << 1, left, right);
113        }
114        // Query right child if needed
115        if (right > mid) {
116            result += query(nodeIndex << 1 | 1, left, right);
117        }
118      
119        return result;
120    }
121  
122    /**
123     * Update parent node's sum based on its children
124     * @param nodeIndex - current node index
125     */
126    private void pushup(int nodeIndex) {
127        tree[nodeIndex].sum = tree[nodeIndex << 1].sum + tree[nodeIndex << 1 | 1].sum;
128    }
129  
130    /**
131     * Push down lazy updates to children nodes
132     * @param nodeIndex - current node index
133     */
134    private void pushdown(int nodeIndex) {
135        if (tree[nodeIndex].lazy == 1) {
136            int mid = (tree[nodeIndex].left + tree[nodeIndex].right) >> 1;
137          
138            // Update left child
139            tree[nodeIndex << 1].sum = mid - tree[nodeIndex].left + 1 - tree[nodeIndex << 1].sum;
140            tree[nodeIndex << 1].lazy ^= 1;
141          
142            // Update right child
143            tree[nodeIndex << 1 | 1].sum = tree[nodeIndex].right - mid - tree[nodeIndex << 1 | 1].sum;
144            tree[nodeIndex << 1 | 1].lazy ^= 1;
145          
146            // Clear current node's lazy flag
147            tree[nodeIndex].lazy ^= 1;
148        }
149    }
150}
151
152/**
153 * Solution class for handling queries on two arrays
154 */
155class Solution {
156    /**
157     * Process queries on two arrays
158     * @param nums1 - first array (binary array)
159     * @param nums2 - second array (integer array)
160     * @param queries - array of queries with 3 types
161     * @return array of results for type 3 queries
162     */
163    public long[] handleQuery(int[] nums1, int[] nums2, int[][] queries) {
164        // Initialize segment tree for nums1
165        SegmentTree tree = new SegmentTree(nums1);
166      
167        // Calculate initial sum of nums2
168        long sum = 0;
169        for (int value : nums2) {
170            sum += value;
171        }
172      
173        // Count type 3 queries to determine result array size
174        int resultCount = 0;
175        for (int[] query : queries) {
176            if (query[0] == 3) {
177                ++resultCount;
178            }
179        }
180      
181        // Initialize result array
182        long[] result = new long[resultCount];
183        int resultIndex = 0;
184      
185        // Process each query
186        for (int[] query : queries) {
187            if (query[0] == 1) {
188                // Type 1: Flip bits in nums1 from index query[1] to query[2]
189                tree.modify(1, query[1] + 1, query[2] + 1);  // Convert to 1-indexed
190            } else if (query[0] == 2) {
191                // Type 2: Add query[1] * (count of 1s in nums1) to each element in nums2
192                sum += 1L * query[1] * tree.query(1, 1, nums2.length);
193            } else {
194                // Type 3: Return current sum of nums2
195                result[resultIndex++] = sum;
196            }
197        }
198      
199        return result;
200    }
201}
202
1class Node {
2public:
3    int left = 0, right = 0;      // Range boundaries [left, right]
4    int sum = 0;                   // Sum of 1s in this range
5    int lazyTag = 0;               // Lazy propagation flag for range flip operations
6};
7
8class SegmentTree {
9public:
10    SegmentTree(vector<int>& nums) {
11        this->nums = nums;
12        int n = nums.size();
13        // Allocate 4n nodes for the segment tree
14        tree.resize(n << 2);
15        for (int i = 0; i < tree.size(); ++i) {
16            tree[i] = new Node();
17        }
18        // Build the segment tree starting from root (index 1)
19        build(1, 1, n);
20    }
21
22    // Range modification: flip all bits in range [l, r]
23    void modify(int nodeIdx, int l, int r) {
24        // If current node's range is completely within [l, r]
25        if (tree[nodeIdx]->left >= l && tree[nodeIdx]->right <= r) {
26            // Toggle lazy flag
27            tree[nodeIdx]->lazyTag ^= 1;
28            // Flip the sum: total_elements - current_sum
29            tree[nodeIdx]->sum = tree[nodeIdx]->right - tree[nodeIdx]->left + 1 - tree[nodeIdx]->sum;
30            return;
31        }
32      
33        // Push down lazy propagation to children
34        pushDown(nodeIdx);
35      
36        int mid = (tree[nodeIdx]->left + tree[nodeIdx]->right) >> 1;
37        // Recursively modify left child if needed
38        if (l <= mid) {
39            modify(nodeIdx << 1, l, r);
40        }
41        // Recursively modify right child if needed
42        if (r > mid) {
43            modify(nodeIdx << 1 | 1, l, r);
44        }
45      
46        // Update current node's sum from children
47        pushUp(nodeIdx);
48    }
49
50    // Query the sum of elements in range [l, r]
51    int query(int nodeIdx, int l, int r) {
52        // If current node's range is completely within [l, r]
53        if (tree[nodeIdx]->left >= l && tree[nodeIdx]->right <= r) {
54            return tree[nodeIdx]->sum;
55        }
56      
57        // Push down lazy propagation to children
58        pushDown(nodeIdx);
59      
60        int mid = (tree[nodeIdx]->left + tree[nodeIdx]->right) >> 1;
61        int result = 0;
62      
63        // Query left child if needed
64        if (l <= mid) {
65            result += query(nodeIdx << 1, l, r);
66        }
67        // Query right child if needed
68        if (r > mid) {
69            result += query(nodeIdx << 1 | 1, l, r);
70        }
71      
72        return result;
73    }
74
75private:
76    vector<Node*> tree;     // Segment tree nodes
77    vector<int> nums;       // Original array
78  
79    // Build the segment tree recursively
80    void build(int nodeIdx, int l, int r) {
81        tree[nodeIdx]->left = l;
82        tree[nodeIdx]->right = r;
83      
84        // Leaf node: directly assign value from nums
85        if (l == r) {
86            tree[nodeIdx]->sum = nums[l - 1];  // Convert to 0-indexed
87            return;
88        }
89      
90        int mid = (l + r) >> 1;
91        // Build left subtree
92        build(nodeIdx << 1, l, mid);
93        // Build right subtree
94        build(nodeIdx << 1 | 1, mid + 1, r);
95      
96        // Update current node's sum from children
97        pushUp(nodeIdx);
98    }
99  
100    // Update parent node's sum from its children
101    void pushUp(int nodeIdx) {
102        tree[nodeIdx]->sum = tree[nodeIdx << 1]->sum + tree[nodeIdx << 1 | 1]->sum;
103    }
104  
105    // Push down lazy propagation to children
106    void pushDown(int nodeIdx) {
107        if (tree[nodeIdx]->lazyTag) {
108            int mid = (tree[nodeIdx]->left + tree[nodeIdx]->right) >> 1;
109          
110            // Apply flip operation to left child
111            tree[nodeIdx << 1]->sum = mid - tree[nodeIdx]->left + 1 - tree[nodeIdx << 1]->sum;
112            tree[nodeIdx << 1]->lazyTag ^= 1;
113          
114            // Apply flip operation to right child
115            tree[nodeIdx << 1 | 1]->sum = tree[nodeIdx]->right - mid - tree[nodeIdx << 1 | 1]->sum;
116            tree[nodeIdx << 1 | 1]->lazyTag ^= 1;
117          
118            // Clear the lazy tag
119            tree[nodeIdx]->lazyTag ^= 1;
120        }
121    }
122};
123
124class Solution {
125public:
126    vector<long long> handleQuery(vector<int>& nums1, vector<int>& nums2, vector<vector<int>>& queries) {
127        // Initialize segment tree with nums1
128        SegmentTree* segTree = new SegmentTree(nums1);
129      
130        // Calculate initial sum of nums2
131        long long totalSum = 0;
132        for (int& value : nums2) {
133            totalSum += value;
134        }
135      
136        vector<long long> answer;
137      
138        // Process each query
139        for (auto& query : queries) {
140            if (query[0] == 1) {
141                // Type 1: Flip bits in nums1 from index query[1] to query[2]
142                // Convert to 1-indexed for segment tree
143                segTree->modify(1, query[1] + 1, query[2] + 1);
144            } 
145            else if (query[0] == 2) {
146                // Type 2: Add query[1] * sum(nums1) to all elements in nums2
147                // This effectively adds query[1] * sum(nums1) to totalSum
148                totalSum += 1LL * query[1] * segTree->query(1, 1, nums1.size());
149            } 
150            else {
151                // Type 3: Return current sum of nums2
152                answer.push_back(totalSum);
153            }
154        }
155      
156        return answer;
157    }
158};
159
1interface Node {
2    left: number;      // Range boundaries [left, right]
3    right: number;
4    sum: number;       // Sum of 1s in this range
5    lazyTag: number;   // Lazy propagation flag for range flip operations
6}
7
8let tree: Node[] = [];
9let nums: number[] = [];
10
11// Build the segment tree recursively
12function build(nodeIdx: number, l: number, r: number): void {
13    tree[nodeIdx].left = l;
14    tree[nodeIdx].right = r;
15  
16    // Leaf node: directly assign value from nums
17    if (l === r) {
18        tree[nodeIdx].sum = nums[l - 1];  // Convert to 0-indexed
19        return;
20    }
21  
22    const mid = (l + r) >> 1;
23    // Build left subtree
24    build(nodeIdx << 1, l, mid);
25    // Build right subtree
26    build((nodeIdx << 1) | 1, mid + 1, r);
27  
28    // Update current node's sum from children
29    pushUp(nodeIdx);
30}
31
32// Update parent node's sum from its children
33function pushUp(nodeIdx: number): void {
34    tree[nodeIdx].sum = tree[nodeIdx << 1].sum + tree[(nodeIdx << 1) | 1].sum;
35}
36
37// Push down lazy propagation to children
38function pushDown(nodeIdx: number): void {
39    if (tree[nodeIdx].lazyTag) {
40        const mid = (tree[nodeIdx].left + tree[nodeIdx].right) >> 1;
41      
42        // Apply flip operation to left child
43        tree[nodeIdx << 1].sum = mid - tree[nodeIdx].left + 1 - tree[nodeIdx << 1].sum;
44        tree[nodeIdx << 1].lazyTag ^= 1;
45      
46        // Apply flip operation to right child
47        tree[(nodeIdx << 1) | 1].sum = tree[nodeIdx].right - mid - tree[(nodeIdx << 1) | 1].sum;
48        tree[(nodeIdx << 1) | 1].lazyTag ^= 1;
49      
50        // Clear the lazy tag
51        tree[nodeIdx].lazyTag ^= 1;
52    }
53}
54
55// Range modification: flip all bits in range [l, r]
56function modify(nodeIdx: number, l: number, r: number): void {
57    // If current node's range is completely within [l, r]
58    if (tree[nodeIdx].left >= l && tree[nodeIdx].right <= r) {
59        // Toggle lazy flag
60        tree[nodeIdx].lazyTag ^= 1;
61        // Flip the sum: total_elements - current_sum
62        tree[nodeIdx].sum = tree[nodeIdx].right - tree[nodeIdx].left + 1 - tree[nodeIdx].sum;
63        return;
64    }
65  
66    // Push down lazy propagation to children
67    pushDown(nodeIdx);
68  
69    const mid = (tree[nodeIdx].left + tree[nodeIdx].right) >> 1;
70    // Recursively modify left child if needed
71    if (l <= mid) {
72        modify(nodeIdx << 1, l, r);
73    }
74    // Recursively modify right child if needed
75    if (r > mid) {
76        modify((nodeIdx << 1) | 1, l, r);
77    }
78  
79    // Update current node's sum from children
80    pushUp(nodeIdx);
81}
82
83// Query the sum of elements in range [l, r]
84function query(nodeIdx: number, l: number, r: number): number {
85    // If current node's range is completely within [l, r]
86    if (tree[nodeIdx].left >= l && tree[nodeIdx].right <= r) {
87        return tree[nodeIdx].sum;
88    }
89  
90    // Push down lazy propagation to children
91    pushDown(nodeIdx);
92  
93    const mid = (tree[nodeIdx].left + tree[nodeIdx].right) >> 1;
94    let result = 0;
95  
96    // Query left child if needed
97    if (l <= mid) {
98        result += query(nodeIdx << 1, l, r);
99    }
100    // Query right child if needed
101    if (r > mid) {
102        result += query((nodeIdx << 1) | 1, l, r);
103    }
104  
105    return result;
106}
107
108// Initialize segment tree with given array
109function initializeSegmentTree(inputNums: number[]): void {
110    nums = inputNums;
111    const n = nums.length;
112    // Allocate 4n nodes for the segment tree
113    tree = new Array(n << 2);
114    for (let i = 0; i < tree.length; i++) {
115        tree[i] = {
116            left: 0,
117            right: 0,
118            sum: 0,
119            lazyTag: 0
120        };
121    }
122    // Build the segment tree starting from root (index 1)
123    build(1, 1, n);
124}
125
126function handleQuery(nums1: number[], nums2: number[], queries: number[][]): number[] {
127    // Initialize segment tree with nums1
128    initializeSegmentTree(nums1);
129  
130    // Calculate initial sum of nums2
131    let totalSum: number = 0;
132    for (const value of nums2) {
133        totalSum += value;
134    }
135  
136    const answer: number[] = [];
137  
138    // Process each query
139    for (const queryItem of queries) {
140        if (queryItem[0] === 1) {
141            // Type 1: Flip bits in nums1 from index query[1] to query[2]
142            // Convert to 1-indexed for segment tree
143            modify(1, queryItem[1] + 1, queryItem[2] + 1);
144        } else if (queryItem[0] === 2) {
145            // Type 2: Add query[1] * sum(nums1) to all elements in nums2
146            // This effectively adds query[1] * sum(nums1) to totalSum
147            totalSum += queryItem[1] * query(1, 1, nums1.length);
148        } else {
149            // Type 3: Return current sum of nums2
150            answer.push(totalSum);
151        }
152    }
153  
154    return answer;
155}
156

Time and Space Complexity

Time Complexity: O(n + m × log n)

The time complexity breaks down as follows:

  • Building the segment tree in the constructor takes O(n) time, where n is the length of nums1. The build method recursively constructs the tree by visiting each node once, and there are O(n) nodes in total.
  • Processing m queries takes O(m × log n) time, where m is the number of queries. For each query:
    • Type 1 queries (modify): Takes O(log n) time as it traverses from root to leaves, visiting at most O(log n) nodes due to the tree height being O(log n).
    • Type 2 queries (query): Takes O(log n) time for the same reason - traversing at most O(log n) nodes to compute the sum.
    • Type 3 queries: Takes O(1) time as it just appends the current sum to the result.
  • The initial sum calculation sum(nums2) takes O(n) time.

Therefore, the total time complexity is O(n + m × log n).

Space Complexity: O(n)

The space complexity consists of:

  • The segment tree array self.tr contains n << 2 (which is 4n) nodes, giving O(n) space.
  • Each node stores a constant amount of data (four integers: l, r, s, lazy).
  • The result array ans stores at most m values, but since we're typically concerned with the space used by the data structure itself rather than the output, the dominant factor is the segment tree.
  • Other variables use O(1) space.

Therefore, the total space complexity is O(n).

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Index Conversion Error (0-indexed vs 1-indexed)

One of the most common pitfalls in this solution is confusion between 0-indexed arrays and 1-indexed segment tree operations. The segment tree internally uses 1-indexed positions while the input arrays nums1 and nums2 are 0-indexed.

Problem Example:

# Incorrect: Forgetting to convert indices
tree.modify(1, l, r)  # Wrong! l and r are 0-indexed from query

# Correct: Converting to 1-indexed
tree.modify(1, l + 1, r + 1)  # Convert 0-indexed to 1-indexed

2. Incorrect Sum Update in Type 2 Query

A critical mistake is misunderstanding how Type 2 queries affect the total sum. The operation adds p * nums1[i] to each nums2[i], which means the total sum increases by p * sum(nums1), NOT p * sum(nums1) * len(nums2).

Problem Example:

# Incorrect: Multiplying by array length
current_sum += param_a * ones_count * len(nums2)  # Wrong!

# Correct: Just multiply by the count of ones
current_sum += param_a * ones_count  # Each 1 in nums1 adds param_a to total

3. Lazy Propagation Flag Management

Failing to properly toggle the lazy flag using XOR can lead to incorrect flip operations. Using assignment instead of XOR will break multiple consecutive flips.

Problem Example:

# Incorrect: Using assignment
self.tree[node_idx].lazy = 1  # Wrong! Doesn't handle multiple flips

# Correct: Using XOR to toggle
self.tree[node_idx].lazy ^= 1  # Correctly handles multiple flips

4. Segment Tree Size Allocation

Allocating insufficient space for the segment tree array is a common mistake. The tree needs up to 4n nodes to handle all possible segments.

Problem Example:

# Incorrect: Insufficient space
self.tree = [Node() for _ in range(n * 2)]  # May cause index out of bounds

# Correct: Allocate 4n space
self.tree = [Node() for _ in range(n << 2)]  # Safe allocation

5. Pushdown Before Query/Modify

Forgetting to push down lazy updates before querying or recursively modifying children will return stale data or apply updates incorrectly.

Problem Example:

def query(self, node_idx: int, left: int, right: int) -> int:
    if self.tree[node_idx].left >= left and self.tree[node_idx].right <= right:
        return self.tree[node_idx].sum
  
    # Incorrect: Missing pushdown before recursive calls
    mid = (self.tree[node_idx].left + self.tree[node_idx].right) >> 1
    result = 0
    # ... recursive calls without pushdown
  
    # Correct: Always pushdown before accessing children
    self.pushdown(node_idx)  # Must come before recursive calls
    mid = (self.tree[node_idx].left + self.tree[node_idx].right) >> 1
    # ... then recursive calls

Solution Best Practices:

  1. Always document index conventions clearly in comments
  2. Test edge cases with single-element arrays and full-range operations
  3. Validate lazy propagation with consecutive flip operations on the same range
  4. Use consistent naming for 0-indexed vs 1-indexed variables
  5. Add assertions during development to catch index errors early
Discover Your Strengths and Weaknesses: Take Our 5-Minute Quiz to Tailor Your Study Plan:

What is an advantages of top-down dynamic programming vs bottom-up dynamic programming?

Recommended Readings

Segment Tree

Segment Tree Faster Range Queries For this article we want to introduce the idea of a Segment Tree Segment Trees allow us to quickly perform range queries as well as range updates Suppose we have an array and we want to know the sum of a particular range of numbers as well

Coding Interview Patterns: The Smart, Data-Driven Way to Prepare for Tech Interviews

Coding Interview Patterns Your Personal Dijkstra's Algorithm to Landing Your Dream Job The goal of AlgoMonster is to help you get a job in the shortest amount of time possible in a data driven way We compiled datasets of tech interview problems and broke them down by patterns This way

Recursion Review

Recursion Recursion is one of the most important concepts in computer science Simply speaking recursion is the process of a function calling itself Using a real life analogy imagine a scenario where you invite your friends to lunch https assets algo monster recursion jpg You first call Ben and ask

Want a Structured Path to Master System Design Too? Don’t Miss This!

Load More