LC-6358.更新数组后处理求和查询

题目描述

leetcode 困难题

给你两个下标从 0 开始的数组 nums1 和 nums2 ，和一个二维数组 queries 表示一些操作。总共有 3 种类型的操作：

• 操作类型 1 为 queries[i] = [1, l, r] 。你需要将 nums1 从下标 l 到下标 r 的所有 0 反转成 1 或将 1 反转成 0 。l 和 r 下标都从 0 开始。
• 操作类型 2 为 queries[i] = [2, p, 0] 。对于 0 <= i < n 中的所有下标，令 nums2[i] = nums2[i] + nums1[i] * p 。
• 操作类型 3 为 queries[i] = [3, 0, 0] 。求 nums2 中所有元素的和。
  请你返回一个数组，包含所有第三种操作类型的答案。

示例1：

输入：nums1 = [1,0,1], nums2 = [0,0,0], queries = [[1,1,1],[2,1,0],[3,0,0]]
输出：[3]
解释：第一个操作后 nums1 变为 [1,1,1] 。第二个操作后，nums2 变成 [1,1,1] ，所以第三个操作的答案为 3 。所以返回 [3] 。

提示1：

1 <= nums1.length,nums2.length <= 10^5
nums1.length = nums2.length
1 <= queries.length <= 10^5
queries[i].length = 3
0 <= l <= r <= nums1.length - 1
0 <= p <= 10^6
0 <= nums1[i] <= 1
0 <= nums2[i] <= 10^9

线段树

显然为区间更新、区间查询问题，可以使用线段树解决。

我们使用线段树维护 $nums1$，而操作 2 只是将 $nums2$ 的数组总和累加一个 $nums1$ 的区间查询，操作 3 只是获取 $nums2$ 的数组总和，所以对于 $nums2$ ，我们只需要使用一个变量维护其数组的总和即可。

注意到操作 1 每次只会将 $nums1$ 某个区间的 $0、1$ 翻转，设区间为 $[i, j]$ ，翻转前区间和为 $k$，那么显然翻转后区间和为 $j - i + 1 - k$ ，也就是说对于线段树的 $update$ 实现，只需要将 $arr[p]$ 更新为 $(j - i + 1) - arr[p]$ 即可。而对于同一个区间的偶数次翻转，又等同于当前，所以我们在进行懒惰标记下放时只需要再判断懒惰标记的值是否为偶数即可。

class Solution {
    public long[] handleQuery(int[] nums1, int[] nums2, int[][] queries) {
        SegTree segTree1 = new SegTree(nums1);
        long sum = Arrays.stream(nums2).mapToLong(Long::valueOf).sum();
        List<Long> ans = new ArrayList<>();
        for (int[] item : queries) {
            if(item[0] == 1){
                segTree1.update(item[1], item[2], 1, 0, nums1.length - 1, 1);
            }
            if (item[0] == 2) {
                sum += segTree1.query(0, nums1.length - 1, 0, nums1.length - 1, 1) * (long)item[1];
            }
            if (item[0] == 3) {
                ans.add(sum);
            }
        }
        return ans.stream().mapToLong(a -> a).toArray();
    }
}

class SegTree {
    long[] arr;
    int[] lazyTag;
    int[] org;

    public SegTree(int[] org) {
        this.org = org;
        this.arr = new long[org.length * 4];
        this.lazyTag = new int[org.length * 4];
        build(0, org.length - 1, 1);
    }

    private void build(int s, int t, int p) {
        if (s == t) {
            arr[p] = org[s];
            return;
        }
        int mid = s + (t - s >> 1);
        build(s, mid, p * 2);
        build(mid + 1, t, p * 2 + 1);
        arr[p] = arr[p * 2] + arr[p * 2 + 1];
    }

    public void update(int l, int r, int c, int s, int t, int p){
        if (l <= s && r >= t) {
            arr[p] = (t - s + 1) - arr[p];
            lazyTag[p] += c;
            return;
        }
        int mid = s + (t - s >> 1);
        if (lazyTag[p] != 0) {
            if (lazyTag[p] % 2 != 0) {
                arr[p * 2] = (mid - s + 1) - arr[p * 2];
            }
            if (lazyTag[p] % 2 != 0){
                arr[p * 2 + 1] = (t - mid) - arr[p * 2 + 1];
            }
            lazyTag[p * 2] += lazyTag[p];
            lazyTag[p * 2 + 1] += lazyTag[p];
            lazyTag[p] = 0;
        }
        if (l <= mid) {
            update(l, r, c, s, mid, p * 2);
        }
        if (r >= mid + 1) {
            update(l, r, c, mid + 1, t, p * 2 + 1);
        }
        arr[p] = arr[p * 2] + arr[p * 2 + 1];
    }

    public long query(int l, int r, int s, int t, int p){
        if (l <= s && r >= t) {
            return arr[p];
        }
        int mid = s + (t - s >> 1);
        if (lazyTag[p] != 0) {
            if (lazyTag[p] % 2 != 0) {
                arr[p * 2] = (mid - s + 1) - arr[p * 2];
            }
            if (lazyTag[p] % 2 != 0){
                arr[p * 2 + 1] = (t - mid) - arr[p * 2 + 1];
            }
            lazyTag[p * 2] += lazyTag[p];
            lazyTag[p * 2 + 1] += lazyTag[p];
            lazyTag[p] = 0;
        }
        int sum = 0;
        if (l <= mid) {
            sum += query(l, r, s, mid, p * 2);
        }
        if (r >= mid + 1) {
            sum += query(l, r, mid + 1, t, p * 2 + 1);
        }
        return sum;
    }

}