E - Count Arithmetic Subarrays 解説 by en_translator
The condition is always satisfied if \(l=r\), so we will add \(n\) to the answer finally, and consider pairs \((l,r)\) such that \(l < r\) from now on.
Let \(B=(B_1,B_2,\dots,B_{N-1})\) be the sequence of differences of \(B=(B_1,B_2,\dots,B_{N-1})\); that is, \(B_i =A_{i+1}-A_i\). Then, \((l,r)\) satisfies the condition, i.e. \(A_l, A_{l+1},\dots,A_r\) is an arithmetic progression, if and only if \(B_l = B_{l+1} = \dots = B_{r-1}\). Thus, the problem can be boiled down as follows:
- Given a sequence \(B=(B_1,B_2,\dots,B_{N-1})\), find the number of pairs \((l, r)\ (1\leq l\leq r\leq N-1)\) such that \(B_l = B_{l+1} = \dots = B_r\).
To solve this problem, we first apply run-length compression to \(B\). That is, we find the sequence of integer pairs \((C_1, D_1), (C_2,D_2),\dots,(C_k, D_k)\) that satisfies the following conditions. This costs \(O(N)\) time.
- Let \(f(i)\) be the sequence where \(C_i\) is repeated \(D_i\) times. \(B\) is a concatenation of \(f(1), f(2), \dots, f(k)\) in this order.
- \(C_i \neq C_{i+1}\) for all \(1\leq i < k\).
Here, for all \(1\leq i < k \), if a segment \([l, r]\) strides over the boundary of \(f(i)\) and \(f(i+1)\) (in other words, if \(l \leq \sum_{j=1}^{i} f(j) < r\)), then that pair \((l, r)\) is not applicable. Thus, a conforming segment \([l, r]\) is completely contained in any of \(f(i9\). Conversely, if the segment \([l, r]\) is completely contained in \(f(i)\), then that pair \((l, r)\) is applicable.
The number of segments completely contained in \(f(i)\) is \(\frac{D_i(D_i+1)}{2}\), so the answer to this problem is \(\displaystyle \sum_{i=1}^{k} \frac{D_i(D_i+1)}{2}\).
Sample code (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;
}
投稿日時:
最終更新: