2569. 更新数组后处理求和查询

2569. 更新数组后处理求和查询

方法：线段树

前置知识

线段树：
线段树 - 多组图带你从头到尾彻底理解线段树

• 线段树节点包含两个属性：区间 和 值；
  • 区间，线段树中的每个节点代表一个区间，即将一个整区间逐级分割为左右两个子区间；
  • 值，当前节点（区间）在题目中所表示的含义。
• 线段树通常包含以下方法，构建树、更新区间、查询区间；
• $lazy-tag$，表示对于某个区间更改后，不立即更新其子节点，而是打上标记，等下次使用到时再进行更新，提高效率。

解题思路

• 对于操作2，由于题中表明 $nums1$ 数组中的值只有 0 / 1，因此每次 $nums2$ 数组中增加的大小实际上是 $nums1$ 数组中当前 1 的个数 $cnt\ast p$；
• 由上述可知，对于某一段区间不需要知道其具体位置上是 0 / 1，只需要维护该区间内 1 的个数（通过 $lazy-\mathrm{线}\mathrm{段}\mathrm{树}$），因此对于操作1，只需要改变区间内 1 的个数；
• 每当遇到操作3，将当前的 $sum$ 值加入答案数组即可。

代码

class Solution {
private:
    static const int MX = 4e5 + 1; // 线段树节点 4n 
    int cnt_one[MX]; // 存储节点[l, r]中1的个数
    bool tag[MX]; // 当前节点打上标记之前，即 tag[i] 为 false 时，不更新其左右子节点的区间
public:
    // 维护节点中1的个数
    void maintain(int h) {
        cnt_one[h] = cnt_one[h * 2] + cnt_one[h * 2 + 1];
    }

    // 构建线段树
    void create_tree(vector<int>&a, int h, int l, int r) {
        if (l == r) {
            cnt_one[h] = a[l - 1];
            return ;
        }
        int m = l + r >> 1;
        create_tree(a, 2 * h, l, m);
        create_tree(a, 2 * h + 1, m + 1, r);
        maintain(h);
    }

    // 执行翻转操作
    void do_work(int h, int l, int r) {
        cnt_one[h] = r - l + 1 - cnt_one[h];
        tag[h] = !tag[h];
    }

    // 更新线段树信息
    // [l, r] 表示线段树节点，[L, R] 表示需要修改的区间
    void update(int h, int l, int r, int L, int R) { 
        if (L <= l && r <= R) {
            do_work(h, l, r);
            return ;
        }
        int m = l + r >> 1;
        if (tag[h]) { // 处理子节点
            do_work(h * 2, l, m);
            do_work(h * 2 + 1, m + 1, r);
            tag[h] = false;
        }
        if (m >= L) update(h * 2, l, m, L, R); // 需要更新的区间有一部分在左子树
        if (m < R) update(h * 2 + 1, m + 1, r, L, R); // ~一部分在右子树
        maintain(h); // 更新后维护当前节点信息
    }

    vector<long long> handleQuery(vector<int>& nums1, vector<int>& nums2, vector<vector<int>>& queries) {
        int n = nums1.size();
        create_tree(nums1, 1, 1, n);
        vector<long long> ans;
        long long sum = accumulate(nums2.begin(), nums2.end(), 0LL);
        for (auto &q : queries) {
            if (q[0] == 1) update(1, 1, n, q[1] + 1, q[2] + 1);
            else if (q[0] == 2) sum += 1LL * q[1] * cnt_one[1];
            else ans.push_back(sum);
        }
        return ans;
    }
};

复杂度分析

时间复杂度：$O(nlogn)$；
空间复杂度：$O(n)$。