けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 166 E - This Message Will Self-Destruct in 5s (1Q, 緑色, 500 点)

上手に式変形しよう!

問題概要

正の整数からなるサイズ  N の数列  A_{1}, A_{2}, \dots, A_{N} が与えられる。次の条件を満たす組  (i, j) の個数を求めよ。

  •  0 \le i \lt j \lt N
  •  A_{i} + A_{j} = j - i

制約

  •  2 \le N \le 2 \times 10^{5}
  •  1 \le A_{i} \le 10^{9}

考えたこと

条件が不思議だ。普通は  i, j よりも  A_{i}, A_{j} の方が大きいことが多いのに、 A_{i} + A_{j} = j - i となる条件を問うとは。

さて、この手の数式を見たときに考えるべきことは「左辺は  i だけ、右辺は  j だけというように式変形する」ということだ。次のようになる。

 A_{i} = A_{j} = j - i
 A_{i} + i = j - A_{i}

この式が示すことは、次のようなことだ。


 j = 0, 1, \dots, N-1 に対して、 B = j - A_{i} としたとき、 A_{i} = B を満たす  i (\lt j) の個数を合算していけばよい


この処理は map などを用いてできる。計算量は  O(N \log N) となる。

コード

#include <bits/stdc++.h>
using namespace std;

// Ai + Aj = j - i
// Ai + i = j - Aj
int main() {
    int N;
    cin >> N;
    vector<long long> A(N);
    for (int i = 0; i < N; i++) cin >> A[i];

    map<long long, long long> Aplus;
    for (int i = 0; i < N; i++) Aplus[A[i] + i]++;
    long long res = 0;
    for (int j = 0; j < N; j++) {
        res += Aplus[j - A[j]];
    }
    cout << res << endl;
}