E - Count Arithmetic Subarrays 解説
by
yuto1115
解説
\(l=r\) のときは常に条件を満たすので、最初に答えに \(n\) を加算しておくことにして、以降は \(l < r\) なる組 \((l,r)\) についてのみ考えます。
\(A\) の差分配列を \(B=(B_1,B_2,\dots,B_{N-1})\) とします。つまり、\(B_i =A_{i+1}-A_i\) です。このとき、\((l,r)\) が条件を満たすこと、すなわち \(A_l, A_{l+1},\dots,A_r\) が等差数列であることは、\(B_l = B_{l+1} = \dots = B_{r-1}\) であることと同値です。よって、本問題は以下のような問題に帰着されます。
- 数列 \(B=(B_1,B_2,\dots,B_{N-1})\) が与えられる。\(B_l = B_{l+1} = \dots = B_r\) であるような組 \((l, r)\ (1\leq l\leq r\leq N-1)\) の数を求めよ。
この問題を解くために、まずは \(B\) をランレングス圧縮します。すなわち、以下の条件を全て満たすような整数の組の列 \((C_1, D_1), (C_2,D_2),\dots,(C_k, D_k)\) を求めます。この操作は \(O(N)\) で行えます。
- \(C_i\) を \(D_i\) 個繋げた列を \(f(i)\) とおく。\(B\) は、\(f(1), f(2), \dots, f(k)\) をこの順にすべて繋げた列である。
- 全ての \(1\leq i < k\) について、\(C_i \neq C_{i+1}\)
ここで、任意の \(1\leq i < k \) について、\(f(i)\) と \(f(i+1)\) の境目をある区間 \([l, r]\) がまたぐ場合(すなわち、\(l \leq \sum_{j=1}^{i} f(j) < r\) である場合)、その組 \((l, r)\) は条件を満たしません。よって、条件を満たす区間 \([l, r]\) はいずれかの \(f(i)\) の中に完全に含まれています。逆に、ある \(f(i)\) の中に区間 \([l, r]\) が完全に含まれている場合、その組 \((l, r)\) は条件を満たします。
\(f(i)\) の中に完全に含まれる区間の数は \(\frac{D_i(D_i+1)}{2}\) ですから、本問題に対する答えは \(\displaystyle \sum_{i=1}^{k} \frac{D_i(D_i+1)}{2}\) です。
実装例 (C++) :
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
ll f(ll n) { return n * (n + 1) / 2; }
int main() {
int n;
cin >> n;
vector<int> a(n);
for (int &i: a) cin >> i;
ll ans = n;
int pre = 0;
for (int i = 1; i < n - 1; i++) {
if (a[i] - a[i - 1] != a[i + 1] - a[i]) {
ans += f(i - pre);
pre = i;
}
}
ans += f(n - 1 - pre);
cout << ans << endl;
}
投稿日時:
最終更新: