[AtCoder] ABC 166 E – This Message Will Self-Destruct in 5s

2020年12月14日

問題

方針

配列の添え字を \( i, j (i < j) \) とすると、

\begin{eqnarray}
j – i &=& A_i + A_j \\
i + A_i &=& j – A_j
\end{eqnarray}

となります。左辺と右辺は、\( i , j \) に関して独立に計算できるので、\( i + A_i \) の頻度を \( B(i + A_i) \) とし、\( i – A_i \) の頻度を \( C(i – A_i) \) とすると、答えは、

\[ \sum_{i = 1}^{N-1}B(i)C(i)\]

となります。

コード

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
int main() {
    int N;
    cin >> N;
    int A[N];
    for (int i = 0; i < N; i++) {
        cin >> A[i];
    }
    ll B[N]{};
    ll C[N]{};
    ll ans = 0;
    for (int i = 0; i < N; i++) {
        if (i + A[i] < N) {
            B[i + A[i]]++;
        }
        if (i - A[i] > 0) {
            C[i - A[i]]++;
        }
    }
    for (int i = 1; i < N - 1; i++) {
        ans += B[i] * C[i];
    }
    cout << ans << "\n";
    return 0;
}

感想

思いつかなかった。