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

题目情境#

题目链接：https://www.luogu.com.cn/problem/P3368 ↗

题目描述

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

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

输入格式

第一行包含两个整数 $N$ 、 $M$ ，分别表示该数列数字的个数和操作的总个数。

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

接下来 $M$ 行每行包含 $2$ 或 $4$ 个整数，表示一个操作，具体如下：

操作 $1$ ：格式：1 x y k 含义：将区间 $[x,y]$ 内每个数加上 $k$ ；

操作 $2$ ：格式：2 x 含义：输出第 $x$ 个数的值。

输出格式

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

输入 #1

plaintext

输出 #1

6
10

plaintext

对于 $100\%$ 的数据： $1 \leq N, M\le 500000$ ， $1 \leq x, y \leq n$ ，保证任意时刻序列中任意元素的绝对值都不大于 $2^{30}$ 。

算法解析：#

区间修改、单点查询可以用树状数组维护差分数组 $d$ ： $d_1=a_1$ ， $d_i=a_i-a_{i-1}$ 。对原数组区间 $[l,r]$ 加上 $k$ ，等价于在差分数组上做两个单点修改 $d_l += k$ 、 $d_{r+1} -= k$ ；查询第 $x$ 个数的值就是差分数组的前缀和 $a_x=\sum_{i=1}^{x}d_i$ 。树状数组恰好支持 $O(\log n)$ 的单点修改和前缀和查询，所以总复杂度为 $O((n+m)\log n)$ 。

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()))
    # 用差分数组初始化：d[1]=a1, d[i]=a[i]-a[i-1]
    prev = 0
    for i, val in enumerate(a, 1):  # 下标从1开始
        d = val - prev
        T.change(i, d)
        prev = val

    for _ in range(m):
        op, *rest = map(int, input().split())
        if op == 1:  # 区间修改： op=1 l r k
            l, r, k = rest
            T.change(l, k)
            if r + 1 <= n:
                T.change(r + 1, -k)
        else:        # 点查询： op=2 x
            x = rest[0]
            print(T.query(x))

python

C++代码实现#

// 树状数组 区修+点查 O(nlogn)
#include<bits/stdc++.h>
using namespace std;

const int N=500010;
int n,m,a[N];

struct BIT{
  int s[N]; //差分的区间和
  void change(int x,int w){
    for(;x<=n;x+=x&-x) s[x]+=w;
  }
  int query(int x){
    int res=0;
    for(;x;x-=x&-x) res+=s[x];
    return res;
  }
}T;
int main(){
  ios::sync_with_stdio(0);
  cin>>n>>m; int op,x,y,k;
  for(int i=1;i<=n;i++) cin>>a[i];
  for(int i=1;i<=m;i++){
    cin>>op>>x;
    if(op==1){
      cin>>y>>k;
      T.change(x,k);
      T.change(y+1,-k); //差分
    }
    else cout<<T.query(x)+a[x]<<"\n";
  }
}