树状数组训练题 - 无聊的数列

P1438 无聊的数列

题目背景

无聊的 YYB 总喜欢搞出一些正常人无法搞出的东西。有一天，无聊的 YYB 想出了一道无聊的题：无聊的数列。。。

题目描述

维护一个数列 \(a_i\)，支持两种操作：

• 1 l r K D：给出一个长度等于 \(r-l+1\) 的等差数列，首项为 \(K\)，公差为 \(D\)，并将它对应加到 \([l,r]\) 范围中的每一个数上。即：令 \(a_l=a_l+K,a_{l+1}=a_{l+1}+K+D\ldots a_r=a_r+K+(r-l) \times D\)。

• 2 p：询问序列的第 \(p\) 个数的值 \(a_p\)。

输入格式

第一行两个整数 \(n,m\) 表示数列长度和操作个数。

第二行 \(n\) 个整数，第 \(i\) 个数表示 \(a_i\)。

接下来的 \(m\) 行，每行先输入一个整数 \(opt\)。

若 \(opt=1\) 则再输入四个整数 \(l,r,K,D\)；

若 \(opt=2\) 则再输入一个整数 \(p\)。

输出格式

对于每个询问，一行一个整数表示答案。

输入输出样例 #1

输入 #1

5 2
1 2 3 4 5
1 2 4 1 2
2 3

输出 #1

6

说明/提示

数据规模与约定

对于 \(100\%\) 数据，\(1\le n,m \le 10^5,-200\le a_i,K,D\le 200, 1 \leq l \leq r \leq n, 1 \leq p \leq n\)。

思路

维护原数组

在未知这道题目是树状数组训练题的时候，我们很容易想到暴力做法。

while (m -- )
{
    int opt;
    cin >> opt;
    if (opt == 1)
    {
        int l, r, k, d;
        cin >> l >> r >> k >> d;
        for (int i = l; i <= r; ++i)
            a[i] += k + (i - l) * 1LL * d;
    }
    else // opt == 2
    {
        int p; cin >> p;
        cout << a[p] << '\n';
    }
}

很容易发现，时间复杂度中，单次修改最坏为\(O(n)\)，查询固定\(O(1)\)，总复杂度最坏\(O(m * n)\)
对于n和m，1e5的数据量会直接炸掉，需要优化

维护斜率和截距

我们可以注意到，每次我们要输出的都是a[p]的最终值，而a[p]是如何一步步地到达最终值我们其实是不关心的。

而且为了减少更改次数，我们需要尽可能地使用常数去描述。很容易发现，等差数列有两个常数，首项\(k\)和公差\(d\)那么这个时候我们就要想到，在修改的操作中要尽可能和这两个数有关。
由于一次函数的集合构成一个线性空间，所以一次函数可以叠加

i	1	2	3	4	5	6
\(a_i\)	4	6	3	5	2	7
\(f(i) = 2i - 4, i \in [3, 5]\)			\(+2*i-4\)	\(+2*i-4\)	\(+2*i-4\)
\(g(i) = 3i + 2, i \in [2, 3]\)		\(+3*i+2\)	\(+3*i+2\)
\(h(i) = -i - 5, i \in [3, 6]\)			\(-1*i-5\)	\(-1*i-5\)	\(-1*i-5\)	\(-1*i-5\)
\(a_i + f(i) + g(i) + h(i)\)	\(+0\)	\(+3*i+2\)	\(+4*i+1\)	\(+1*i-9\)	\(+1*i-9\)	\(-1*i-5\)

如图，我们需要的就是最后一行的东西。这样我们就会很自然地想到，去开两个数组来维护斜率和截距的变化了。同时，每次修改都是对区间加上一个常数，这也就引起了我们往后向树状数组或者线段树的方向去思考。

while (m--)
{
    int opt;
    cin >> opt;
    if (opt == 1)
    {
        int l, r, k, d;
        cin >> l >> r >> k >> d;
        i64 A = k - 1LL * d * l;   // 本次操作的常数项
        for (int i = l; i <= r; ++i)
        {
            bb[i] += A;
            kk[i] += d;
        }
    }
    else
    {
        int p;
        cin >> p;
        i64 ans = a[p] + bb[p] + kk[p] * p;
        cout << ans << '\n';
    }
}

树状数组优化

显然刚刚的操作时间复杂度也是不对的。

看到给定数据为1e5很容易想到\(n \dot \log{n}\)的算法。那么单次操作就要\(\log{n}\)看到\(\log{n}\)级别的，区间修改和查询，第一反应其实是线段树。但是这里差分+树状数组已经足够。

于是就有了完整的AC代码：

#include <bits/stdc++.h>
using namespace std;
using i32 = int;
using i64 = long long;

template<typename T>
struct BIT
{
    T n; vector<T> tr;

    BIT(T n_)
    {
        n = n_;
        tr.assign(n + 1, 0);
    }

    inline T lowbit(T x)
    {
        return x & -x;
    }

    void add(T idx, T delta)
    {
        while (idx <= n)
        {
            tr[idx] += delta;
            idx += lowbit(idx);
        }
    }

    T pref(T idx)
    {
        T ans = 0;
        while (idx > 0)
        {
            ans += tr[idx];
            idx -= lowbit(idx);
        }
        return ans;
    }
};

i32 main()
{
    ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
    i32 n, m; cin >> n >> m;
    vector<i64> arr(n + 1);

    for (int i = 1; i <= n; ++ i)
        cin >> arr[i];

    BIT<i64> bit_d = BIT<i64>(n); // 斜率d
    BIT<i64> bit_k = BIT<i64>(n); // 截距(k - d * l)

    while (m -- )
    {
        i32 opt; cin >> opt;
        if (opt == 1)
        {
            i32 l, r, k, d; cin >> l >> r >> k >> d;

            bit_d.add(l, d);
            bit_d.add(r + 1, -d);

            bit_k.add(l, k - 1LL * d * l);
            bit_k.add(r + 1, - (k - 1LL * d * l));
        }

        if (opt == 2)
        {
            i32 p; cin >> p;
            cout << arr[p] + bit_d.pref(p) * p + bit_k.pref(p) << '\n';
        }
    }


    return 0;
}