Skip to content

2569. Handling Sum Queries After Update

Description

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

  1. 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.
  2. 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.
  3. 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 <= 105
  • nums1.length = nums2.length
  • 1 <= queries.length <= 105
  • 0 <= nums1[i] <= 1
  • 0 <= nums2[i] <= 109

 

Solutions

Solution: Segment Tree

  • Time complexity: O(logn)
  • Space complexity: O(4n -> n)

 

JavaScript

js
/**
 * @param {number[]} nums1
 * @param {number[]} nums2
 * @param {number[][]} queries
 * @return {number[]}
 */
const handleQuery = function (nums1, nums2, queries) {
  const tree = new SegmentTree(nums1);
  const result = [];
  let sum = nums2.reduce((total, num) => total + num);

  for (const [type, l, r] of queries) {
    if (type === 1) {
      tree.updateRange(l, r);
    } else if (type === 2) {
      sum += l * tree.getSum();
    } else {
      result.push(sum);
    }
  }

  return result;
};

class SegmentTree {
  constructor(nums) {
    const n = nums.length;

    this.n = n;
    this.nums = nums;
    this.tree = Array.from({ length: 4 * n }, () => 0);
    this.lazy = Array.from({ length: 4 * n }, () => false);
    this.build(0, 0, n - 1);
  }

  build(index, start, end) {
    if (start === end) {
      this.tree[index] = this.nums[start];
      return;
    }

    const mid = Math.floor((start + end) / 2);

    this.build(index * 2 + 1, start, mid);
    this.build(index * 2 + 2, mid + 1, end);
    this.tree[index] = this.tree[index * 2 + 1] + this.tree[index * 2 + 2];
  }

  updateRange(l, r) {
    return this.#updateRange(0, 0, this.n - 1, l, r);
  }

  #updateRange(index, start, end, l, r) {
    if (this.lazy[index]) {
      this.flip(index, start, end);
      this.lazy[index] = false;
    }

    if (start > r || end < l) return;

    if (start >= l && end <= r) {
      this.flip(index, start, end);
      return;
    }

    const mid = Math.floor((start + end) / 2);

    this.#updateRange(index * 2 + 1, start, mid, l, r);
    this.#updateRange(index * 2 + 2, mid + 1, end, l, r);
    this.tree[index] = this.tree[index * 2 + 1] + this.tree[index * 2 + 2];
  }

  flip(index, start, end) {
    this.tree[index] = end - start + 1 - this.tree[index];

    if (start < end) {
      this.lazy[index * 2 + 1] = !this.lazy[index * 2 + 1];
      this.lazy[index * 2 + 2] = !this.lazy[index * 2 + 2];
    }
  }

  getSum() {
    return this.tree[0];
  }
}

Released under the MIT license