C++的数据结构(八）：线段树

news2024/9/25 21:27:25

线段树是一种高效的树形数据结构，用于处理区间查询和区间更新问题。它的基本思想是将一个大的区间分解为若干个小的、不相交的区间，每个小区间对应线段树中的一个节点。线段树的每个节点保存了该区间的信息（如区间最大值、区间和等），这使得我们能在对数时间内完成区间查询和更新操作。

线段树的构建通常采用递归方式，自底向上逐层构建。每个节点的信息由其子节点的信息合并而来，如区间和节点的值即为其左右子区间节点值的和。如下图所示。

构建线段树时，我们通常需要一个原始数组和一个递归函数。递归函数的作用是确定每个节点所代表的区间范围，并根据需要计算并保存节点的信息。代码如下。

struct Node {
    int start, end;
    int sum; // 假设我们存储的是区间和
    Node *left, *right;
    Node(int s, int e) : start(s), end(e), sum(0), left(nullptr), right(nullptr) {}
};

Node* build(int arr[], int start, int end) {
    if (start > end) return nullptr;
    Node* root = new Node(start, end);
    if (start == end) {
        root->sum = arr[start];
        return root;
    }
    int mid = (start + end) / 2;
    root->left = build(arr, start, mid);
    root->right = build(arr, mid + 1, end);
    root->sum = root->left->sum + root->right->sum;
    return root;
}

区间查询是线段树的一个重要应用，可以快速获取某个区间内的信息。查询时，我们根据查询区间的位置与线段树节点的区间关系，决定是递归查询左子树、右子树，还是将两者的信息合并。代码如下。

int query(Node* root, int qs, int qe) {
    if (root->start > qe || root->end < qs) return 0;
    if (qs <= root->start && qe >= root->end) return root->sum;
    int sum = 0;
    sum += query(root->left, qs, qe);
    sum += query(root->right, qs, qe);
    return sum;
}

区间更新操作通常涉及对某个区间内的元素进行加减操作或其他更复杂的操作。区间更新需要同时维护一个懒惰标记（Lazy Propagation）数组，用于记录节点下尚未传导到子节点的更新信息。懒惰标记处理的代码通常较为复杂，示例为比较简单的代码。

void update(Node* root, int index, int value) {
    // 如果当前节点的区间不包含index，直接返回
    if (root->start > index || root->end < index) return;

    // 如果当前节点是一个叶子节点，直接更新值
    if (root->start == root->end) {
        root->sum = value;
        return;
    }

    // 递归更新左子树或右子树
    update(root->left, index, value);
    update(root->right, index, value);

    // 递归回溯时更新当前节点的值
    root->sum = root->left->sum + root->right->sum;
}

下面是一个完整的例子，展示如何使用线段树解决区间和查询与更新的问题。代码如下。

#include <iostream>
#include <vector>
using namespace std;
struct Node {
    int start, end;
    int sum; // 假设我们存储的是区间和
    Node *left, *right;
    Node(int s, int e) : start(s), end(e), sum(0), left(nullptr), right(nullptr) {}
};

Node* build(int arr[], int start, int end) {
    if (start > end) return nullptr;
    Node* root = new Node(start, end);
    if (start == end) {
        root->sum = arr[start];
        return root;
    }
    int mid = (start + end) / 2;
    root->left = build(arr, start, mid);
    root->right = build(arr, mid + 1, end);
    root->sum = root->left->sum + root->right->sum;
    return root;
}
int query(Node* root, int qs, int qe) {
    if (root->start > qe || root->end < qs) return 0;
    if (qs <= root->start && qe >= root->end) return root->sum;
    int sum = 0;
    sum += query(root->left, qs, qe);
    sum += query(root->right, qs, qe);
    return sum;
}

void update(Node* root, int index, int value) {
    // 如果当前节点的区间不包含index，直接返回
    if (root->start > index || root->end < index) return;

    // 如果当前节点是一个叶子节点，直接更新值
    if (root->start == root->end) {
        root->sum = value;
        return;
    }

    // 递归更新左子树或右子树
    update(root->left, index, value);
    update(root->right, index, value);

    // 递归回溯时更新当前节点的值
    root->sum = root->left->sum + root->right->sum;
}

int main(){
	 int arr[] = {1, 3, 5, 7, 9, 11};
    int n = sizeof(arr) / sizeof(arr[0]);
    Node* root = build(arr, 0, n - 1);

    // 查询区间 [1, 4] 的和
    cout << "Sum from index 1 to 4: " << query(root, 1, 4) << endl; 

    // 更新 
    update(root, 2, 10);

    // 再次查询区间 [1, 4] 的和，应该会反映更新
    cout << "Sum from index 1 to 4 after update: " << query(root, 1, 4) << endl; 
	
	return 0;
}

结果如下图所示。