树状数组的区间修改与区间查询

S(x) = d[1] 
     + (d[1] + d[2]) 
     + (d[1] + d[2] + d[3]) 
     + ... 
     + (d[1] + d[2] + ... + d[x])

S(x) = x * d[1] + (x-1) * d[2] + ... + 1 * d[x]

S(x) = (x+1) * (d[1] + d[2] + ... + d[x]) - (1*d[1] + 2*d[2] + ... + x*d[x])

S(x) = (x+1) * sum(d[1..x]) - sum(i * d[i] for i=1..x)

prefixSum(x) = (x+1) * bit1.sum(x) - bit2.sum(x)

区间和 [l, r] = prefixSum(r) - prefixSum(l-1)

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

class FenwickRange {
private:
    int n;
    vector<long long> bit1; // 维护差分数组 d[i] 的前缀和
    vector<long long> bit2; // 维护 i * d[i] 的前缀和

    // lowbit 运算，树状数组基础
    long long lowbit(long long x) { return x & -x; }

    // 单点更新某个树状数组（bit）
    void add(vector<long long>& bit, int idx, long long delta) {
        while (idx <= n) {
            bit[idx] += delta;
            idx += lowbit(idx);
        }
    }

    // 查询某个树状数组的前缀和（1..idx）
    long long sum(const vector<long long>& bit, int idx) {
        long long res = 0;
        while (idx > 0) {
            res += bit[idx];
            idx -= lowbit(idx);
        }
        return res;
    }

public:
    // 构造函数：初始化大小为 n，下标从1开始
    FenwickRange(int n) : n(n), bit1(n + 1, 0), bit2(n + 1, 0) {}

    // 区间修改：给区间 [l, r] 加上 k
    void rangeAdd(int l, int r, long long k) {
        add(bit1, l, k);          // d[l] += k
        add(bit1, r + 1, -k);     // d[r+1] -= k
        add(bit2, l, l * k);      // l*d[l] 增加 l*k
        add(bit2, r + 1, -(r + 1) * k); // (r+1)*d[r+1] 减少 (r+1)*k
    }

    // 求原数组的前缀和 a[1] + ... + a[x]
    long long prefixSum(int x) {
        if (x <= 0) return 0;
        return (x + 1) * sum(bit1, x) - sum(bit2, x);
    }

    // 区间查询：求 a[l] + ... + a[r]
    long long rangeSum(int l, int r) {
        return prefixSum(r) - prefixSum(l - 1);
    }
};

int main() {
    FenwickRange ft(10); // 创建管理下标1..10的树状数组
    // 初始所有值均为0
    ft.rangeAdd(1, 5, 1);     // 第1~5位加1
    ft.rangeAdd(3, 7, 2);     // 第3~7位加2
    ft.rangeAdd(2, 4, -1);    // 第2~4位减1

    cout << "前缀和[1..3]: " << ft.prefixSum(3) << endl;   // 应输出3
    cout << "区间和[2..6]: " << ft.rangeSum(2, 6) << endl; // 应输出9

    // 手动验证各位置最终值
    // 位置: 1  2  3  4  5  6  7  8 9 10
    // 值:   1  0  2  2  3  2  2  0 0  0
    cout << "位置4的值: " << ft.rangeSum(4, 4) << endl; // 应输出2
    return 0;
}

class FenwickRange:
    def __init__(self, n: int):
        self.n = n
        self.bit1 = [0] * (n + 1)  # 维护差分数组 d[i] 的前缀和
        self.bit2 = [0] * (n + 1)  # 维护 i * d[i] 的前缀和

    def lowbit(self, x: int) -> int:
        return x & -x

    def _add(self, bit: list, idx: int, delta: int) -> None:
        """单点更新某个树状数组"""
        while idx <= self.n:
            bit[idx] += delta
            idx += self.lowbit(idx)

    def _sum(self, bit: list, idx: int) -> int:
        """查询某个树状数组的前缀和"""
        res = 0
        while idx > 0:
            res += bit[idx]
            idx -= self.lowbit(idx)
        return res

    def range_add(self, l: int, r: int, k: int) -> None:
        """区间修改：[l, r] 加上 k"""
        self._add(self.bit1, l, k)          # d[l] += k
        self._add(self.bit1, r + 1, -k)     # d[r+1] -= k
        self._add(self.bit2, l, l * k)      # l*d[l] 增加 l*k
        self._add(self.bit2, r + 1, -(r + 1) * k)  # (r+1)*d[r+1] 减少 (r+1)*k

    def prefix_sum(self, x: int) -> int:
        """求原数组的前缀和 a[1] + ... + a[x]"""
        if x <= 0:
            return 0
        return (x + 1) * self._sum(self.bit1, x) - self._sum(self.bit2, x)

    def range_sum(self, l: int, r: int) -> int:
        """区间查询：[l, r] 的和"""
        return self.prefix_sum(r) - self.prefix_sum(l - 1)


if __name__ == "__main__":
    ft = FenwickRange(10)  # 下标1..10，初始全0
    ft.range_add(1, 5, 1)     # 第1~5位加1
    ft.range_add(3, 7, 2)     # 第3~7位加2
    ft.range_add(2, 4, -1)    # 第2~4位减1

    print("前缀和[1..3]:", ft.prefix_sum(3))   # 输出 3
    print("区间和[2..6]:", ft.range_sum(2, 6))  # 输出 9