官方题解｜ 统计操作后不同的序列

题目解析

每次操作后都可以得到一个序列，$(L, R)$ 二元组的总数量是 $\dfrac{N(N + 1)}{2}$。

我们考虑将其中操作后得到的重复序列去除掉。

令 $X = \dfrac{\sum_{i=L}^R A_i}{R - L + 1}$，那么对于 $A_L = X$ 或 $A_R = X$ 的情况，我们将 $L$ 向右边移动 $1$ 或将 $R$ 向左移动 $1$，操作后的序列不会发生变化。

我们考虑找到那些使得 $(A_L, \ldots, A_R)$ 的平均值与 $A_L$ 或 $A_R$ 相等的 $(L, R)$ 将其从总数中去掉。

具体地，对于 $x = 1, \ldots, \max{A}$，进行如下操作：
定义 $(B_1, \ldots, B_N) = (A_1 - x, \ldots, A_N - x)$ 并构造 $B$ 的前缀和序列 $(S_0 = 0, S_1 = B_1, \ldots, S_N = \sum_{i=1}^N B_i)$。

对于 $i = 1, \ldots, N$：

• 若 $B_i = 0$，则从答案中减去满足 $0 \le j < i$ 且 $S_j = S_i$ 的 $j$ 的个数。
  对应平均值为 $x$、右端为 $x$ 而左端任意的情况。

• 若 $B_i \ne 0$，则从答案中减去满足 $0 \le j < i$ 且 $B_{j+1} = 0$ 且 $S_j = S_i$ 的 $j$ 的个数。
  对应于平均值为 $x$、右端不为 $x$ 而左端为 $x$ 的情况。

对于每个 $x$ 的计算时间复杂度为 $\mathrm{O}(N)$。

因此，整体的时间复杂度为 $\mathrm{O}(N\max{A})$。

AC 代码

#include <bits/stdc++.h>

using i64 = long long;

constexpr int M = 30;

int main() {
    int n; std::cin >> n;
    std::vector<int> a(n + 1);
    for (int i = 1; i <= n; ++i)
        std::cin >> a[i];
    i64 res = n * (n + 1LL) / 2; 
    for (int x = 1; x <= M; ++x) {
        std::vector<int> b(n + 1), s(n + 1);
        for (int i = 1; i <= n; ++i)
            b[i] = a[i] - x;
        std::partial_sum(b.begin(), b.end(), s.begin());
        const int DELTA = n * M;
        auto f = [&](int v) {return DELTA + v;};
        std::vector<int> zero(DELTA * 2 + 1);
        std::vector<int> nzero(DELTA * 2 + 1);
        for (int i = 1; i <= n; ++i)
            if (b[i] == 0) {
                zero[f(s[i - 1])] += 1;
                res -= zero[f(s[i])] + nzero[f(s[i])];
            }
            else {
                nzero[f(s[i - 1])] += 1;
                res -= zero[f(s[i])];
            }
    }
    std::cout << res + 1 << "\n";
    return 0;
}