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\) と等しいものをすべて見つけることに言い換えることができます。
このことを考えると、連想配列などを使った次のようなアルゴリズムが考えられます。
- 整数をキーとして非負整数を返す連想配列 \(\operatorname{map}\) を作り、はじめすべて \(0\) で初期化する。
- \(\operatorname{ans}\) を \(0\) で初期化する。
- \(j=2,3,\ldots,N\) に対して、次を繰り返す。
- \(\operatorname{ans}\leftarrow\operatorname{ans}+\operatorname{map}\lbrack j-A _ j\rbrack\) と更新する。
- \(\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:
