【AtCoder】ABC 417 C - Distance Indicators | 緑コーダーが解くAtCoder

atcoder.jp

配点: 300 点 / 実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: 246 / NoviSteps: 2Q

問題概要
制約
考察
実装例
コメント

問題概要

長さ $N$ の整数列 $A$ が与えられる。2整数の組 $(i,j) \: (1 \leq i \lt j \leq N)$ のうち、 $j-i=A_i+A_j$ を満たすものの個数を求めよ。

制約

$1\le N\le2\times10^5$
$1\le A_i\le2\times10^5$
入力はすべて整数

考察

例によって $(i, j)$ の組を愚直に調べていくことはできない。

ここで、条件式を

$\begin{align*} j - i = A_i + A_j \iff j - A_j = A_i + i \end{align*}$

のように、添字をそろえるように変形してみる。

すると、各 $j \: (j = 1, 2, \dots, N)$ に対して、 $1 \leq i \lt j$ を満たす $i$ について、 $j - A_j = A_i + i$ を満足する個数を求めればよいことになる。

この個数は、あらかじめmapで $A_i + i$ の個数を記録しておくことで高速に求めることができる。

全体の計算量は $O(N \log N)$ となり、本問の制約下では十分高速。

実装例

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

#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)

using ll = long long;

// ======================================== //

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

    map<int, int> mp;
    ll ans = 0;
    rep(j, 0, N)
    {
        int A;
        cin >> A;

        ans += mp[j - A];
        mp[A + j]++;
    }

    cout << ans << endl;

    return 0;
}