C++实现线段树

迈克柯里喵

于 2024-07-16 17:25:54 发布

阅读量171

点赞数 3

文章标签： c++ 算法

本文链接：https://blog.csdn.net/weixin_44286126/article/details/140472131

版权

📺 教程

推荐视频讲解：线段树+懒标记

🧠 记忆

idx 从1开始数，数组区间也是从1开始数，比如 [1,4] 表示第1个数字到第4个数字，所以为了便于写代码，nums数组建议从下标1开始存数据，0位空着不写；
build 只手动执行一遍，我们只需要调用 query、update；
无论是 build query 还是 update，idx 参数始终填 1；
build需要pushup；query向下递归需要pushdown；update向下递归前要pushdown，递归完成后要pushup

build 只会执行一次，build 遇到叶子节点，则把数组值存到sum里，否则继续build左右区间：

void build(int idx, int left, int right) {
    tr[idx] = {left,right,0,0};
    if (left == right) {  // 是叶子节点
        tr[idx].sum = nums[left];  // 叶子节点存储数组值
        return;  // 是叶子节点，则返回
    }

    // 不是叶子节点，则裂开
    int mid = (right - left) / 2 + left;
    build(2*idx, left, mid);
    build(2*idx+1, mid+1, right);
    pushup(idx);  // 向上更新
}

build完叶子节点后，需要向上更新，获得父节点的区间和：

void pushup(int idx) {  // 向上更新，根据两个孩子，计算父节点的区间和
    tr[idx].sum = tr[2*idx].sum + tr[2*idx+1].sum;
}

区间修改，给区间 [x,y]，即第x到第y个数每个数字加 k，此处使用 🚩懒标记，一旦发现xy覆盖了节点对应的区间，则更新该节点的sum值，然后懒标记加上k，就返回；如果没有覆盖，则继续向下递归，向下递归前需要先向下更新：

void update(int idx, int x, int y, int k) {  // 区间修改
    if (x <= tr[idx].left && tr[idx].right <= y) {  // 覆盖则修改
        tr[idx].sum += (tr[idx].right - tr[idx].left + 1) * k;
        tr[idx].add +=k;
        return;
    }

    // 不覆盖则裂开
    pushdown(idx);
    int mid = (tr[idx].right - tr[idx].left) / 2 + tr[idx].left;
    if (x <= mid) update(2*idx, x, y, k);
    if (y > mid) update(2*idx+1, x, y, k);
    pushup(idx);
}

向下更新：

void pushdown(int idx) {  // 向下更新
    if (tr[idx].add != 0) {
        int left  = 2*idx;
        int right = 2*idx+1;
        tr[left].sum  += tr[idx].add * (tr[left].right  - tr[left].left  + 1);
        tr[right].sum += tr[idx].add * (tr[right].right - tr[right].left + 1);
        tr[left].add  += tr[idx].add;
        tr[right].add += tr[idx].add;
        tr[idx].add = 0;
    }
}

区间查询，返回节点的区间和。

int query(int idx, int x, int y) {  // 区间查询
    if (x <= tr[idx].left && tr[idx].right <= y) {  // 覆盖则返回
        return tr[idx].sum;
    }

    // 不覆盖则裂开
    pushdown(idx);
    int mid = (tr[idx].right - tr[idx].left) / 2 + tr[idx].left;
    int sum = 0;
    if (x <= mid) sum += query(2*idx, x, y);
    if (y > mid)  sum += query(2*idx+1, x, y);
    return sum;
}

✍ 练习

练习题1 P3374：

// https://www.luogu.com.cn/problem/P3374

#include <iostream>
#include <string>
#include <vector>
#include <sstream>
using namespace std;

struct node {
    int left;    // 左边界
    int right;    // 右边界
    int sum;  // 区间和
    int add;  // 懒标记
};

int n, m;  // n,m
vector<int> nums;  // nums
vector<node> tr;
vector<vector<int>> op;

void pushup(int idx) {  // 向上更新，根据两个孩子，计算父亲的区间和
    tr[idx].sum = tr[2*idx].sum + tr[2*idx+1].sum;
}

void pushdown(int idx) {  // 向下更新
    if (tr[idx].add != 0) {
        int left  = 2*idx;
        int right = 2*idx+1;
        tr[left].sum  += tr[idx].add * (tr[left].right  - tr[left].left  + 1);
        tr[right].sum += tr[idx].add * (tr[right].right - tr[right].left + 1);
        tr[left].add  += tr[idx].add;
        tr[right].add += tr[idx].add;
        tr[idx].add = 0;
    }
}

void build(int idx, int left, int right) {
    tr[idx] = {left,right,0,0};
    if (left == right) {
        tr[idx].sum = nums[left];  // 
        return;  // 是叶子节点，则返回
    }

    // 不是叶子节点，则裂开
    int mid = (right - left) / 2 + left;
    build(2*idx, left, mid);
    build(2*idx+1, mid+1, right);
    pushup(idx);
}

void update(int idx, int x, int y, int k) {  // 区间修改
    if (x <= tr[idx].left && tr[idx].right <= y) {  // 覆盖则修改
        tr[idx].sum += (tr[idx].right - tr[idx].left + 1) * k;
        tr[idx].add +=k;
        return;
    }

    // 不覆盖则裂开
    pushdown(idx);
    int mid = (tr[idx].right - tr[idx].left) / 2 + tr[idx].left;
    if (x <= mid) update(2*idx, x, y, k);
    if (y > mid) update(2*idx+1, x, y, k);
    pushup(idx);
}

int query(int idx, int x, int y) {  // 区间查询
    if (x <= tr[idx].left && tr[idx].right <= y) {  // 覆盖则返回
        return tr[idx].sum;
    }

    // 不覆盖则裂开
    pushdown(idx);
    int mid = (tr[idx].right - tr[idx].left) / 2 + tr[idx].left;
    int sum = 0;
    if (x <= mid) sum += query(2*idx, x, y);
    if (y > mid)  sum += query(2*idx+1, x, y);
    return sum;
}

int main() {
    // input
    string input;
    getline(cin, input);
    stringstream ss(input);
    ss >> n >> m;

    nums.resize(n+1);  // nums
    tr.resize((n+1)*4);  // tr
    string input1;
    getline(cin, input1);
    stringstream ss1(input1);
    for (int i = 1; i <= n; i++) {
        ss1 >> nums[i];
    }
    
    vector<vector<int>> _op(m, vector<int>(3));  // 操作
    string input2;
    for (int i = 0; i < m; i++) {
        getline(cin, input2);
        stringstream ss2(input2);
        ss2 >> _op[i][0] >> _op[i][1] >> _op[i][2];
    }
    op = _op;
    // end input
    
    build(1, 1, n);
    for (int i = 0; i < m; i++) {
        if (op[i][0] == 1) {
            update(1, op[i][1], op[i][1], op[i][2]);
        }
        else if (op[i][0] == 2) {
            cout << query(1, op[i][1], op[i][2]) << endl;
        }
    }
    return 0;
}

练习题2 P3372：

// https://www.luogu.com.cn/problem/P3372

#include <bits/stdc++.h>
using namespace std;

struct node {
    int left;    // 左边界
    int right;    // 右边界
    long long sum;  // 区间和
    long long add;  // 懒标记
};

int n, m;  // n,m
vector<int> nums;  // nums
vector<node> tr;
vector<vector<int>> op;

void pushup(int idx) {  // 向上更新，根据两个孩子，计算父亲的区间和
    tr[idx].sum = tr[2*idx].sum + tr[2*idx+1].sum;
}

void pushdown(int idx) {  // 向下更新
    if (tr[idx].add != 0) {
        int left  = 2*idx;
        int right = 2*idx+1;
        tr[left].sum  += tr[idx].add * (tr[left].right  - tr[left].left  + 1);
        tr[right].sum += tr[idx].add * (tr[right].right - tr[right].left + 1);
        tr[left].add  += tr[idx].add;
        tr[right].add += tr[idx].add;
        tr[idx].add = 0;
    }
}

void build(int idx, int left, int right) {
    tr[idx] = {left,right,0,0};
    if (left == right) {
        tr[idx].sum = nums[left];
        return;  // 是叶子节点，则返回
    }

    // 不是叶子节点，则裂开
    int mid = (right - left) / 2 + left;
    build(2*idx, left, mid);
    build(2*idx+1, mid+1, right);
    pushup(idx);
}

void update(int idx, int x, int y, int k) {  // 区间修改
    if (x <= tr[idx].left && tr[idx].right <= y) {  // 覆盖则修改
        tr[idx].sum += (tr[idx].right - tr[idx].left + 1) * k;
        tr[idx].add += k;
        return;
    }

    // 不覆盖则裂开
    pushdown(idx);
    int mid = (tr[idx].right - tr[idx].left) / 2 + tr[idx].left;
    if (x <= mid) update(2*idx, x, y, k);
    if (y > mid) update(2*idx+1, x, y, k);
    pushup(idx);
}

long long query(int idx, int x, int y) {  // 区间查询
    if (x <= tr[idx].left && tr[idx].right <= y) {  // 覆盖则返回
        return tr[idx].sum;
    }

    // 不覆盖则裂开
    pushdown(idx);
    int mid = (tr[idx].right - tr[idx].left) / 2 + tr[idx].left;
    long long sum = 0;
    if (x <= mid) sum += query(2*idx, x, y);
    if (y > mid)  sum += query(2*idx+1, x, y);
    return sum;
}

int main() {
    // input
    string input;
    getline(cin, input);
    stringstream ss(input);
    ss >> n >> m;

    nums.resize(n+1);  // nums
    tr.resize((n+1)*4);  // tr
    string input1;
    getline(cin, input1);
    stringstream ss1(input1);
    for (int i = 1; i <= n; i++) {
        ss1 >> nums[i];
    }
    
    vector<vector<int>> _op(m, vector<int>(4));  // 操作
    string input2;
    for (int i = 0; i < m; i++) {
        getline(cin, input2);
        stringstream ss2(input2);
        ss2 >> _op[i][0] >> _op[i][1] >> _op[i][2];
        if (_op[i][0] == 1) ss2 >> _op[i][3];
    }
    op = _op;
    // end input
    
    build(1, 1, n);
    for (int i = 0; i < m; i++) {
        if (op[i][0] == 1) {
            update(1, op[i][1], op[i][2], op[i][3]);
        }
        else if (op[i][0] == 2) {
            cout << query(1, op[i][1], op[i][2]) << endl;
        }
    }
    return 0;
}