【LeetCode 算法】Handling Sum Queries After Update 更新数组后处理求和查询-Segment Tree

Handling Sum Queries After Update 更新数组后处理求和查询

问题描述：

给你两个下标从 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]\ast p$ 。
操作类型 3 为 $queries[i]=[3,0,0]$ 。求 nums2 中所有元素的和。
请你返回一个数组，包含所有第三种操作类型的答案。

$$\begin{gathered} 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 \end{gathered}$$

分析

这个问题给的2个数组nums1,nums2,同时还有query指令。

问题本身比较绕，$nums1$是01数组，而$nums2$可以看成是由$nums1$的一个函数映射。
而每次的$query$指令为3的时候，需要计算出$nums2$ 的元素和。

如果是纯粹计算数组的元素和，是配不上这个hard。
没错，问题的query指令还有1和2，它们会对nums1进行修改，从而影响到nums2的元素，进而会影响到结果。

最原始的方式就是暴力。

• $query==1$的时候，会将$nums1$中的区间$[l,r]$的$01$进行反转，即0变1,1变0，而产生的影响就是$[l,r]$内的出现的1，会导致处于该区间内的$nums2[i]$,增加$nums1[i]\ast p$.

• $query==1$的时候，对$nums1$修改的时间复杂度为$O(k)$,k为区间的范围，

• $query==2$，对$nums2$的修改时间复杂度是$O(N)$。

• $query==3$，要计算$sum(nums2)$,时间复杂度$O(N)$.

那么整体的时间复杂度就是$O(N^2)$,到此暴力的模拟思路就可以完成了，结果TLE。

其实nums2这个数组，是一个中间数组，虽然是统计nums2的和，实际上并不需要得到nums2 的详细元素。

在该问题中nums2的和在指令执行的过程中一定是单调递增的，而且在$query=1$之后，区间$[l,r]$产生了$x个1$，那么这$x个1$会影响到对应$nums2$对应这个区间的$nums2[i]$,效果就是这个区间的nums2的和增加了x*p 。

如果一开始nums1有x个1，那么nums2初始和为$su{m}_1$，当进行了$query==1$ 的操作，nums1的某个区间出现了$x_1$个1，此时nums2的元素和就是$su{m}_1+x_1\ast p$.

所以如果可以快速的统计$nums1$在每个$query==1$的操作下，区间产生的1的数量，就可以快速的累加到nums2的元素和上，这样就可以在$O(1)$的时间复杂度下得到$query==3$时的结果。

而能够达到这个效果的，就是线段树。而且线段树在解决这类问题中，更新和查询的时间复杂度都是$O(logN)$,所以整体的时间复杂度就是$O(NlogN)$

代码

线段树

public long[] handleQuery(int[] nums1, int[] nums2, int[][] queries) {
        int n = nums1.length, m = 0, i = 0;
        cnt1 = new int[n * 4];
        flip = new boolean[n * 4];
        build(nums1, 1, 1, n);

        long sum = 0L;
        for (var x : nums2){
            sum += x;
        }
            
        for (var q : queries){
            if (q[0] == 3) ++m;
        }
            
        long[] ans = new long[m];
        for (var q : queries) {
            if (q[0] == 1){
              	update(1, 1, n, q[1] + 1, q[2] + 1);  
            } 
            else if (q[0] == 2){
	            sum += (long) q[1] * cnt1[1];  
            } 
            else{
                ans[i++] = sum;
            } 
        }
        return ans;
    }
	// segment tree
    private int[] cnt1;
    private boolean[] flip;

    // 维护区间 1 的个数
    private void maintain(int o) {
        cnt1[o] = cnt1[o * 2] + cnt1[o * 2 + 1];
    }

    // 执行区间反转
    private void do_(int o, int l, int r) {
        cnt1[o] = r - l + 1 - cnt1[o];
        flip[o] = !flip[o];
    }

    // 初始化线段树   o,l,r=1,1,n
    private void build(int[] a, int o, int l, int r) {
        if (l == r) {
            cnt1[o] = a[l - 1];
            return;
        }
        int m = (l + r) / 2;
        build(a, o * 2, l, m);
        build(a, o * 2 + 1, m + 1, r);
        maintain(o);
    }

    // 反转区间 [L,R]   o,l,r=1,1,n
    private void update(int o, int l, int r, int L, int R) {
        if (L <= l && r <= R) {
            do_(o, l, r);
            return;
        }
        int m = (l + r) / 2;
        if (flip[o]) {
            do_(o * 2, l, m);
            do_(o * 2 + 1, m + 1, r);
            flip[o] = false;
        }
        if (m >= L){
            update(o * 2, l, m, L, R);
        } 
        if (m < R){
          	update(o * 2 + 1, m + 1, r, L, R);  
        } 
        maintain(o);
    }

时间复杂度$O(NlogN)$

空间复杂度$O(N)$

代码来源于 灵神大佬,略有调整，线段树写的很简洁。

Tag

Array

Segment Tree