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

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

2024-11-23

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;
}