线段树入门：更新数组后处理求和查询

更新数组后处理求和查询

给你两个下标从 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$ 。

示例 2：

输入：nums1 = $1$ , nums2 = $5$ , queries = $\[2,0,0$ , $3,0,0$ ]

输出： $5$

解释：第一个操作后，nums2 保持不变为 $5$ ，所以第二个操作的答案是 5 。所以返回 $5$ 。

解题思路

对于操作2，由于题中表明数组中的值只有，因此每次数组中增加的大小实际上是数组中当前的个数；由上述可知，对于某一段区间不需要知道其具体位置上是，只需要维护该区间内的个数（通过线段树），因此对于操作，只需要改变区间内的个数；每当遇到操作，将当前的值加入答案数组即可。

参考代码

复制代码

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;
    }
};