C1 树状数组-点修区查 • EH's Blog

题目情境#

题目描述

如题，已知一个数列，你需要进行下面两种操作：

将某一个数加上 $x$ ；
求出某区间每一个数的和。

输入格式

第一行包含两个正整数 $n,m$ ，分别表示该数列数字的个数和操作的总个数。

第二行包含 $n$ 个用空格分隔的整数，其中第 $i$ 个数字表示数列第 $i$ 项的初始值。

接下来 $m$ 行每行包含 $3$ 个整数，表示一个操作，具体如下：

1 x k 含义：将第 $x$ 个数加上 $k$ ；
2 x y 含义：输出区间 $[x,y]$ 内每个数的和。

输出格式

输出包含若干行整数，即为所有操作 $2$ 的结果。

输入 #1

plaintext

输出 #1

14
16

plaintext

【数据范围】对于 $100\%$ 的数据， $1\le n,m \le 5\times 10^5$ 。

算法解析：#

树状数组（Fenwick Tree）用 $O(\log n)$ 维护”单点修改 + 区间求和”。核心是 $lowbit(x) = x\ \&\ (-x)$ ，它取 $x$ 二进制最低位的 $1$ 对应的值：change 修改时沿 $x + lowbit(x)$ 向上传播，query 查询时沿 $x - lowbit(x)$ 向下累加前缀和，区间 $[l,r]$ 的和就是 $query(r)-query(l-1)$ 。

Python代码实现#

import sys
input = sys.stdin.readline

class BIT:
    def __init__(self, n):
        self.n = n
        self.s = [0] * (n + 1)

    def change(self, x, w):
        while x <= self.n:
            self.s[x] += w
            x += x & -x

    def query(self, x):
        res = 0
        while x > 0:
            res += self.s[x]
            x -= x & -x
        return res

if __name__ == "__main__":
    n, m = map(int, input().split())
    T = BIT(n)

    a = list(map(int, input().split()))
    for i, k in enumerate(a, 1):  # 下标从1开始
        T.change(i, k)

    for _ in range(m):
        op, *rest = map(int, input().split())
        if op == 1:  # 单点修改
            x, k = rest
            T.change(x, k)
        else:        # 区间查询
            x, y = rest
            print(T.query(y) - T.query(x - 1))

python

import sys
input = sys.stdin.readline
sys.setrecursionlimit(10 ** 6)

class Node():
    def __init__(self, l, r, he):
        self.l = l
        self.r = r
        self.he = he

def build(p, l, r):
    tr[p] = Node(l, r, 0)
    if l == r:
        tr[p].he = w[l]
        return
    m = (l + r) // 2
    build(p * 2, l, m)
    build(p * 2 + 1, m + 1, r)
    tr[p].he = tr[p * 2].he + tr[p * 2 + 1].he

def update(p, idx, k):
    if tr[p].l == idx and tr[p].r == idx:
        tr[p].he += k
        return
    m = (tr[p].l + tr[p].r) // 2
    if idx <= m:
        update(p * 2, idx, k)
    if idx > m:
        update(p * 2 + 1, idx, k)
    tr[p].he = tr[p * 2].he + tr[p * 2 + 1].he

def query(p, l, r):
    if tr[p].l >= l and tr[p].r <= r:
        return tr[p].he
    m = (tr[p].l + tr[p].r) // 2
    he = 0
    if l <= m:
        he += query(p * 2, l, r)
    if r > m:
        he += query(p * 2 + 1, l, r)
    return he

n, m = map(int, input().split())
w = list(map(int, input().split()))
N = n * 4 + 1
tr = [Node(0, 0, 0) for _ in range(N * 4)]
build(1, 0, n - 1)
for i in range(m):
    tmp = list(map(int, input().split()))
    if tmp[0] == 2:
        print(query(1, tmp[1] - 1, tmp[2] - 1))
    else:
        update(1, tmp[1] - 1, tmp[2])

python

C++代码实现#

#include <bits/stdc++.h>
using namespace std;
const int N = 500010;
int n, m;
struct BIT {
    vector<int> s; // 用 vector 更灵活
    BIT(int n = 0) { s.assign(n + 1, 0); }
    void change(int x, int w) {
        for (; x < (int)s.size(); x += x & -x) s[x] += w;
    }
    int query(int x) {
        int res = 0;
        for (; x; x -= x & -x) res += s[x];
        return res;
    }
    int kth(int k) {
        int pos = 0;
        int sum = 0;
        int step = 1;
        while ((step << 1) < (int)s.size())
            step <<= 1;
        for (; step; step >>= 1) {
            int nxt = pos + step;
            if (nxt < (int)s.size() && sum + s[nxt] < k) {
                pos = nxt;
                sum += s[nxt];
            }
        }
        return pos + 1;
    }
};
int main() {
    ios::sync_with_stdio(0);
    cin.tie(0);
    cin >> n >> m;
    BIT T(n);  // 局部创建，大小为 n
    int op, x, y, k;
    for (int i = 1; i <= n; i++) {
        cin >> k;
        T.change(i, k);
    }
    for (int i = 1; i <= m; i++) {
        cin >> op >> x;
        if (op == 1) {
            cin >> k;
            T.change(x, k);
        } else {
            cin >> y;
            cout << T.query(y) - T.query(x - 1) << "\n";
        }
    }
}

上文的kth其实可以不用，那个是树状数组上二分的过程。

它返回：最小的 pos，满足 query(pos) >= k。

另外，有小线段树：

//洛谷P3374
#include<bits/stdc++.h>
using namespace std;
#define endl "\n"
#define range(i, a, b) for (int i = (a); i < (b); ++i)
#define lc (p<<1)
#define rc (p<<1|1)
#define N 500005
int n, w[N];
struct node {
    int l, r, sum;
} tr[N * 4];

void build(int p, int l, int r) {
    tr[p].l = l;
    tr[p].r = r;
    tr[p].sum = w[l];
    if (l == r) return;
    int m = (l + r) / 2;
    build(lc, l, m);
    build(rc, m + 1, r);
    tr[p].sum = tr[lc].sum + tr[rc].sum;
}

// 点修改（从根递归进入）
void update(int p, int x, int k) {
    if (tr[p].l == x && tr[p].r == x) {
        tr[p].sum += k;
        return;
    }
    int m = (tr[p].l + tr[p].r) / 2;
    if (x <= m) update(lc, x, k);
    else update(rc, x, k);
    tr[p].sum = tr[lc].sum + tr[rc].sum;
}

int query(int p, int x, int y) {
    if (x <= tr[p].l && tr[p].r <= y) return tr[p].sum;
    int m = (tr[p].l + tr[p].r) / 2;
    int sum = 0;
    if (x <= m) sum += query(lc, x, y);
    if (y > m)  sum += query(rc, x, y);
    return sum;
}
int main(){
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
    int n, m;
    cin >> n >> m;
    range(i, 0, n){
        cin >> w[i + 1];
    }
    build(1, 1, n);
    range(i, 0, m){
        int op; cin >> op;
        int x, y, k;
        if (op == 1){
            cin >> x >> k;  // 单点更新
            update(1, x, k);
        }
        else{
            cin >> x >> y;  // 区间查询
            cout << query(1, x, y) << endl;
        }
    }
}