Official

C - Distance Indicators Editorial by MMNMM


ありえるすべての \((i,j)\) の組(\(\dfrac{N(N-1)}2\) 通りあります)に関して判定を行うと、実行時間制限に間に合わせるのは難しいです。

ここで、判定する式 \(j-i=A _ i+A _ j\) を変形して、\(j-A _ j=i+A _ i\) としてみましょう。 すると、\(j\) を \(1\) つ固定して条件を満たす \((i,j)\) の組をすべて見つけることを、\(1\le i\lt j\) を満たす \(i\) のうち、\(i+A _ i\) の値が \(j-A _ j\) と等しいものをすべて見つけることに言い換えることができます。

このことを考えると、連想配列などを使った次のようなアルゴリズムが考えられます。

  1. 整数をキーとして非負整数を返す連想配列 \(\operatorname{map}\) を作り、はじめすべて \(0\) で初期化する。
  2. \(\operatorname{ans}\) を \(0\) で初期化する。
  3. \(j=2,3,\ldots,N\) に対して、次を繰り返す。
    1. \(\operatorname{ans}\leftarrow\operatorname{ans}+\operatorname{map}\lbrack j-A _ j\rbrack\) と更新する。
    2. \(\operatorname{map}\lbrack j+A _ j\rbrack\leftarrow\operatorname{map}\lbrack j+A _ j\rbrack+1\) と更新する。

3-2. のステップが行われることによって、3-1. のステップの時点で \(\operatorname{map}\lbrack x\rbrack\) には \(i+A _ i=x\) を満たす \(1\le i\lt j\) の個数が入っているので、このアルゴリズムで正しい答えを得ることができます。

\(1\le A _ i\) なので \(i+A _ i=j-A _ j=x\) を満たす \(x\) は \(0\le x\lt N\) を満たす(もう少し狭めることもできます)ことを利用して連想配列を長さ \(N\) の配列に置き換えることもできます。

実装例は以下のようになります。

答えが \(2 ^ {32}\) を越える場合があることに注意してください(例えば、\(N=2\times10 ^ 5, A _ i=\vert i-10 ^ 5\vert\) の場合)。

#include <iostream>
#include <map>
using namespace std;

int main() {
    int N;
    cin >> N;

    // counter[x] : x = i + A[i] を満たす i の個数 
    map<int, int> counter;

    long ans = 0;
    for (int i = 0; i < N; ++i) {
        int a;
        cin >> a;
        // j - A[j] = i + A[i] を満たす個数を加算して
        ans += counter[i - a];
        // j + A[j] のカウンタを進める
        ++counter[i + a];
    }
    cout << ans << endl;
    return 0;
}

posted:
last update: