线段树模板：区间和与区间更新

线段树是一种非常强大的数据结构，它允许我们在对数时间内对数组的区间执行查询和更新操作。在这个教程中，我们将通过一个具体的例子——使用线段树来解决洛谷P3372 【模板】线段树 1“区间求和”和“区间更新”的问题来学习线段树的基础模板的应用。

我并不打算动动手指就能把线段树的工作原理描述得非常细致易懂。

如果读者完全不明白线段树，建议在拥有二叉树的前置知识下，对着这里的代码敲一遍，相信你会豁然开朗，自己敲一遍比看别人说一千次还要有效。

线段树的基本概念

线段树是一个二叉树，每个节点代表一个区间的聚合信息（例如区间和、最大值或最小值）。对于长度为n的数组，线段树大约需要4 * n的空间来存储（至于为什么是4 * n的空间，请看oi-wiki，这里不多赘述）。

构建线段树

构建线段树涉及将数组分解为多个区间，并将这些区间作为树的叶子节点。内部节点存储子区间的聚合值。构建树的过程是递归的，时间复杂度为O(n)。

构建函数

    inline void buildTree(int l, int r, int node_idx = 1) {
    

        tree[node_idx].l = l;
        tree[node_idx].r = r;

        if (l == r) {
    
            tree[node_idx].sum = nums[l];
            return;
        }

        int double_idx = (node_idx << 1);
        int left_chil_idx = double_idx + 1;
        int right_chil_idx = double_idx + 2;

        int mid = l + ((r - l) >> 1);

        buildTree(l, mid, left_chil_idx);
        buildTree(mid + 1, r, right_chil_idx);

        tree[node_idx].sum = tree[left_chil_idx].sum + tree[right_chil_idx].sum;
    }

在这个函数中，我们从根节点开始，递归地将数组分成两部分，直到每个区间只包含一个元素。我们将数组的元素存储在树的叶子节点，然后计算内部节点的值，它是其左右子节点值的和。

懒惰更新

down函数，将之前保存在lazy数组中没有真正进行更新操作的数字，真正执行一次更新操作。

在区间查询和区间更新操作中，如果访问到了某一个节点，并且还要继续深入访问其子节点，就应该将之前保存下来的懒惰更新的值进行真正的更新。

inline void down(int node_idx) {
    

    if (lazy[node_idx] == 0) return;

    int tree_left = tree[node_idx].l;
    int tree_right = tree[node_idx].r;
    int mid = tree_left + ((tree_right - tree_left) >> 1);

    int double_idx = (node_idx << 1);
    int left_chil_idx = double_idx + 1;
    int right_chil_idx = double_idx + 2;

    tree[left_chil_idx].sum += lazy[node_idx] * (mid - tree_left + 1);
    tree[right_chil_idx].sum += lazy[node_idx] * (tree_right - mid);
    lazy[left_chil_idx] += lazy[node_idx];
    lazy[right_chil_idx] += lazy[node_idx];

    lazy[node_idx] = 0;

}

区间查询

区间查询是线段树的核心功能之一。我们可以使用线段树在O(log n)时间内查询任意区间的聚合信息。

查询函数

inline int queryRange(int query_left, int query_right, int node_idx = 1) {
    

    int tree_left = tree[node_idx].l, tree_right = tree[node_idx].r;

    // 完全不与查询目标的区间重合
    if (tree_left > query_right || tree_right < query_left) return 0;
    // 完全与查询目标的区间重合
    if (tree_left >= query_left && tree_right <= query_right) {
    
        return tree[node_idx].sum;
    }

    down(node_idx);

    // 部分重合的情况下，我们要把查询目标的区间一分为二，分别查询左右两半。
    int double_idx = (node_idx << 1);
    int left_chil_idx = double_idx + 1;
    int right_chil_idx = double_idx + 2;


    int mid = tree_left + ((tree_right - tree_left) >> 1);

    // 这里的mid有三种情况：
    // 1.mid可能刚好位于[query_left, query_right]之间，刚好可以将其分为两半。
    // 2.mid可能比原来的查询目标的右端点要大。
    // 3.mid + 1可能比原来的查询目标的左端点要小。

    // 因此直接写query(query_left, mid, left_chil_idx);的话
    // 有多算一部分区间的风险。

    int left_query = queryRange(query_left, min(query_right, mid), left_chil_idx);
    int right_query = queryRange(max(query_left, mid + 1), query_right, right_chil_idx);

    return left_query + right_query;

}

在查询函数中，我们首先检查当前区间是否与查询区间重叠。如果没有重叠，返回0。如果当前区间完全包含于查询区间中，返回当前节点的值。如果当前区间与查询区间部分重叠，则递归查询左右子树。

注意：

• 因为当[sl, sr]部分重叠目标区间的时候，它的中点mid，可能会出现mid + 1 < l，或者mid > r的情况，这时候分别要进行min(mid, r)和max(mid + 1, l), r)以防止错误修改原来的目标区间。

区间更新

线段树也支持区间更新，我们可以在O(log n)时间内更新区间内所有元素的值。

更新函数

inline void addRange(int query_left, int query_right, int val, int node_idx = 1) {
    
    int tree_left = tree[node_idx].l, tree_right = tree[node_idx].r;

    // 完全不与查询目标的区间重合
    if (tree_left > query_right || tree_right < query_left) return;
    // 完全与查询目标的区间重合
    if (tree_left >= query_left && tree_right <= query_right) {
    
        tree[node_idx].sum += val * (tree_right - tree_left + 1);
        lazy[node_idx] += val; // 懒惰更新，只更新当前节点，打上懒惰更新中待更新的值，以后访问到再来更新。
        return;
    }
    
    // 将以前懒惰更新中记录的值应用到子节点中。
    down(node_idx);

    int double_idx = (node_idx << 1);
    int left_chil_idx = double_idx + 1;
    int right_chil_idx = double_idx + 2;

    int mid = tree_left + ((tree_right - tree_left) >> 1);
    // 这里的min和max的使用的原因请看上面的queryRange函数。
    addRange(query_left, min(query_right, mid), val, left_chil_idx);
    addRange(max(query_left, mid + 1), query_right, val, right_chil_idx);

    tree[node_idx].sum = tree[left_chil_idx].sum + tree[right_chil_idx].sum;
}

区间查询和区间更新时应该注意的点

更新和查询的时候，我们会把区间分为两半，这意味着我们需要把查询目标区间分裂成两个子区间

1. 那么现在可以想象，我们当前节点代表[tree_left, tree_right]这一区间，而我们的查询目标区间是[query_left, query_right]

2. 取mid为中点，该节点的左右子节点分别代表这两个子区间：[tree_left, mid]，[mid + 1, tree_right]，我们要在这两个子节点中进行子查询

3. 由于左右子节点分别代表：[tree_left, mid]和[mid + 1, tree_right]，它们可能会各占[query_left, query_right]一部分，因此它们都有各自的子目标区间。

4. 我们要将query_left, query_right两者与mid进行比较，让左右子节点的查询获取正确的子目标区间的左右端点。

情况解析

对于mid和[query_left, query_right]，有以下三种情况：

情况 1: mid刚好位于[query_left, query_right]之间

tree_left	…	query_left	…	mid	…	query_right	…	tree_right

说明:

• ... 表示区间中的其余部分，长度至少为1，下面的表格同理。
• 此时[query_left, mid]，[mid + 1, query_right]刚好可以将查询目标分为两半而不出错误。

情况 2: mid比原来的查询目标的右端点要大

tree_left	…	query_left	…	query_right	…	mid	…	tree_right

说明:

• 此时[query_left, mid]，[mid + 1, query_right]将会脱离原本的查询目标，导致查询结果偏大。

情况 3: mid + 1比原来的查询目标的左端点要小

tree_left	…	mid + 1	…	query_left	…	query_right	…	tree_right

说明:

• 与第2种情况同理，在这种情况下，查询结果偏大。

以上三种情况的表格展示了mid在不同位置对查询区间的影响。因此代码实现的时候不能直接把mid当成查询目标的区间的左右端点。

使用线段树

以下是如何使用线段树进行区间查询和更新的完整代码示例：

#include <iostream>
#include <vector>

using namespace std;

#define int long long

int n, m;

struct Node {
    
    int l = 0, r = 0, sum = 0;
};

class SegmentTree {
    
public:
    vector<Node> tree;
    vector<int> lazy;
    vector<int> nums;

    inline void buildTree(int l, int r, int node_idx = 1) {
    

        tree[node_idx].l = l;
        tree[node_idx].r = r;

        if (l == r) {
    
            tree[node_idx].sum = nums[l];
            return;
        }

        int double_idx = (node_idx << 1);
        int left_chil_idx = double_idx + 1;
        int right_chil_idx = double_idx + 2;

        int mid = l + ((r - l) >> 1);

        buildTree(l, mid, left_chil_idx);
        buildTree(mid + 1, r, right_chil_idx);

        tree[node_idx].sum = tree[left_chil_idx].sum + tree[right_chil_idx].sum;
    }

    inline void down(int node_idx) {
    

        if (lazy[node_idx] == 0) return;

        int tree_left = tree[node_idx].l;
        int tree_right = tree[node_idx].r;
        int mid = tree_left + ((tree_right - tree_left) >> 1);

        int double_idx = (node_idx << 1);
        int left_chil_idx = double_idx + 1;
        int right_chil_idx = double_idx + 2;

        tree[left_chil_idx].sum += lazy[node_idx] * (mid - tree_left + 1);
        tree[right_chil_idx].sum += lazy[node_idx] * (tree_right - mid);
        lazy[left_chil_idx] += lazy[node_idx];
        lazy[right_chil_idx] += lazy[node_idx];

        lazy[node_idx] = 0;

    }

    inline int queryRange(int query_left, int query_right, int node_idx = 1) {
    

        int tree_left = tree[node_idx].l, tree_right = tree[node_idx].r;

        // 完全不与查询目标的区间重合
        if (tree_left > query_right || tree_right < query_left) return 0;
        // 完全与查询目标的区间重合
        if (tree_left >= query_left && tree_right <= query_right) {
    
            return tree[node_idx].sum;
        }

        down(node_idx);

        // 部分重合的情况下，我们要把查询目标的区间一分为二，分别查询左右两半。
        int double_idx = (node_idx << 1);
        int left_chil_idx = double_idx + 1;
        int right_chil_idx = double_idx + 2;


        int mid = tree_left + ((tree_right - tree_left) >> 1);

        // 这里的mid有三种情况：
        // 1.mid可能刚好位于[query_left, query_right]之间，刚好可以将其分为两半。
        // 2.mid可能比原来的查询目标的右端点要大。
        // 3.mid + 1可能比原来的查询目标的左端点要小。

        // 因此直接写query(query_left, mid, left_chil_idx);的话
        // 有多算一部分区间的风险。

        int left_query = queryRange(query_left, min(query_right, mid), left_chil_idx);
        int right_query = queryRange(max(query_left, mid + 1), query_right, right_chil_idx);

        return left_query + right_query;

    }

    inline void addRange(int query_left, int query_right, int val, int node_idx = 1) {
    
        int tree_left = tree[node_idx].l, tree_right = tree[node_idx].r;

        // 完全不与查询目标的区间重合
        if (tree_left > query_right || tree_right < query_left) return;
        // 完全与查询目标的区间重合
        if (tree_left >= query_left && tree_right <= query_right) {
    
            tree[node_idx].sum += val * (tree_right - tree_left + 1);
            lazy[node_idx] += val; // 懒惰更新，只更新当前节点，打上懒惰更新中待更新的值，以后访问到再来更新。
            return;
        }
        // 将以前懒惰更新中记录的值应用到子节点中。
        down(node_idx);

        int double_idx = (node_idx << 1);
        int left_chil_idx = double_idx + 1;
        int right_chil_idx = double_idx + 2;

        int mid = tree_left + ((tree_right - tree_left) >> 1);
        // 这里的min和max的使用的原因请看上面的queryRange函数。
        addRange(query_left, min(query_right, mid), val, left_chil_idx);
        addRange(max(query_left, mid + 1), query_right, val, right_chil_idx);

        tree[node_idx].sum = tree[left_chil_idx].sum + tree[right_chil_idx].sum;
//        cout << "ADD RANGE HERE" <<endl;
    }

    explicit SegmentTree(const vector<int> &ns) : nums(ns) {
    
        this->tree.resize(4 * n);
        this->lazy.resize(4 * n);
        buildTree(1, n);
    }


};

inline int rd() {
    
    int c = getchar(), f = 1, res = 0;
    while (c > '9' || c < '0') {
    
        if (c == '-') f = -1;
        c = getchar();
    }

    while (c <= '9' && c >= '0') {
    
        res = (res << 3) + (res << 1) + c - '0';
        c = getchar();
    }
    return res * f;
}

signed main() {
    
    n = rd(), m = rd();

    vector<int> nums(n + 1);
    for (int i = 1; i <= n; ++i) {
    
        nums[i] = rd();
    }

    SegmentTree sgt(nums);

    int type, x, y, k;

    while (m--) {
    
        type = rd();
        x = rd(), y = rd();

        if (type == 1) {
    
            k = rd();
            sgt.addRange(x, y, k);
        } else {
    
            printf("%lld\n", sgt.queryRange(x, y));
        }
    }
    return 0;
}

在主函数中，我们首先读取数组长度和操作数，然后使用输入数组构建线段树。接下来，根据操作类型进行查询或更新。

总结

线段树是解决区间查询和更新问题的非常有效的数据结构。理解它的关键点在于理解懒惰更新这一个点，总的来说，这里的线段树其实就是一个完全二叉树，借助懒惰更新实现了更低的复杂度，如果不熟悉线段树就对着代码重新敲一遍，原理很好理解。