C - Distance Indicators Editorial by en_translator
If we try all pairs \((i,j)\) (there are \(\dfrac{N(N-1)}2\) pairs), it is difficult to finish within the execution time limit.
Here, let us deform the equation \(j-i=A _ i+A _ j\) into \(j-A _ j=i+A _ i\). Then, finding all pairs \((i,j)\) for a fixed \(j\) is equivalent to finding the number of indices \(i\) with \(1\le i\lt j\), such that the value \(i+A _ i\) is equal to \(j-A _ j\).
Given this fact, we come up with the following algorithm using an associated array.
- Prepare an associated array \(\operatorname{map}\) that accepts an integer as a key and returns a non-negative integer, initialized with \(0\).
- Initialize \(\operatorname{ans}\) with \(0\).
- For \(j=2,3,\ldots,N\), repeat the following.
- Set \(\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\).
Due to step 3-2., when step 3-1. is performed \(\operatorname{map}\lbrack x\rbrack\) stores the number of indices \(1\le i\lt j\) with \(i+A _ i=x\), so this algorithm yields a correct answer.
Since \(1\le A _ i\), the value \(x\) satisfying \(i+A _ i=j-A _ j=x\) is within \(0\le x\lt N\) (it can be narrowed down even further, by the way). This allows us to replace the associated array with an array of length \(N\).
The following is sample code.
Note that the answer may exceed \(2 ^ {32}\). (For example, when \(N=2\times10 ^ 5\) and \(A _ i=\vert i-10 ^ 5\vert\).)
#include <iostream>
#include <map>
using namespace std;
int main() {
int N;
cin >> N;
// counter[x] : the number of i satisfying x = i + A[i]
map<int, int> counter;
long ans = 0;
for (int i = 0; i < N; ++i) {
int a;
cin >> a;
// Add the count satisfying j - A[j] = i + A[i]
ans += counter[i - a];
// increment the counter for j + A[j]
++counter[i + a];
}
cout << ans << endl;
return 0;
}
posted:
last update: