c++区间修改求和,带惰性计算的实现:
//
// Created by wwww on 2023/5/1.
//
#ifndef LEETCODE_SEGMENT_TREE_HPP
#define LEETCODE_SEGMENT_TREE_HPP
#include <bits/stdc++.h>
using namespace std;
template <typename T> struct BasicSegTreeNode {
using DataType = T;
using p_node = BasicSegTreeNode<T> *;
BasicSegTreeNode(size_t l, size_t r, DataType d)
: data(d), left_bound(l), right_bound(r), right_child(nullptr),
left_child(nullptr), lazy{} {};
BasicSegTreeNode(size_t l, size_t r) : BasicSegTreeNode(l, r, DataType{}){};
void push_down() {
if (!lazy)
return;
left_child->lazy_inc(lazy);
right_child->lazy_inc(lazy);
lazy = DataType{};
}
void pull_up() { data = left_child->data + right_child->data; }
void lazy_inc(DataType inc) {
lazy += inc;
data += inc * (right_bound - left_bound + 1);
}
bool within(size_t ql, size_t qr) const {
return ql <= left_bound && right_bound <= qr;
}
bool overlap(size_t ql, size_t qr) const {
return !(ql > right_bound || qr < left_bound);
}
DataType val() { return data; }
decltype(auto) left() { return *left_child; }
decltype(auto) right() { return *right_child; }
void set_child(BasicSegTreeNode<T> &l, BasicSegTreeNode<T> &r) {
left_child = &l, right_child = &r;
}
private:
p_node left_child, right_child;
size_t left_bound, right_bound;
DataType data, lazy;
};
template <typename T = int> class SegTree {
using Node = BasicSegTreeNode<T>;
vector<Node> tree;
std::reference_wrapper<const vector<T>> arr;
public:
explicit SegTree(const vector<T> &v) : arr(v) {
int n = arr.get().size();
tree.reserve(n << 1);
_build(0, n - 1);
}
SegTree() = delete;
T range_sum(int l, int r) { return _range_query(l, r, tree.front()); }
void range_add(int l, int r, int val) {
_range_modify(l, r, val, tree.front());
}
private:
T _range_query(int ql, int qr, Node &cur) {
if (!cur.overlap(ql, qr))
return {};
if (cur.within(ql, qr))
return cur.val();
cur.push_down();
return _range_query(ql, qr, cur.left()) + _range_query(ql, qr, cur.right());
}
void _range_modify(int ql, int qr, T val, Node &cur) {
if (!cur.overlap(ql, qr))
return;
if (cur.within(ql, qr)) {
cur.lazy_inc(val);
return;
}
cur.push_down();
_range_modify(ql, qr, val, cur.left());
_range_modify(ql, qr, val, cur.right());
cur.pull_up();
}
Node &_build(size_t l, size_t r) {
if (l == r) {
tree.emplace_back(l, r, arr.get()[l]);
return tree.back();
}
tree.emplace_back(l, r);
auto &ret = tree.back();
auto m = l + ((r - l) >> 1);
ret.set_child(_build(l, m), _build(m + 1, r));
ret.pull_up();
return ret;
}
};
#endif // LEETCODE_SEGMENT_TREE_HPP