AtCoder ABC 166 E問題「This Message Will Self-Destruct in 5s」見直し

問題文（意訳）

パーティの参加者はN人で、それぞれ 1から Nまでの番号がついています。
参加者 iの身長は Aiで、「A1,A2,A3,,,,,AN」の様な身長のリストを利用し以下の条件を満たす 2人が何組存在するか求めなさい。
・2人の持つ番号の差の絶対値が、2人の身長の和に等しい。

入力

N
A1 A2 A3 ... AN

私が考えたのは以下の様な感じで、全部の組み合わせのリストを作って「i-j=Ai+Aj」であればカウントアップするという方法

import itertools
N = int(input())
h_list = list(map(int, input().split()))
c_list = list(itertools.combinations(range(N),2))
pair_of_spies = 0

for c in c_list:
       if abs(c[0] - c[1]) == h_list[c[0]] + h_list[c[1]]:
           pair_of_spies += 1

print(pair_of_spies)

当然タイムアウトでした。

入力の最大値は2*10^5らしいので組み合わせのパターンは「19,999,900,000」になります。

200000 choose 2 - Google 検索

組み合わせの計算は上記のようにGoogleで検索できます。

早速解答を見ていきましょう

まず、条件を式にすると「i < j」の時
j - i = Aj + Ai　です。
上記を式変形すると以下の様になります
Ai + i = Aj - j 　
これは「番号と身長を足した値」と、それ以降で「番号から身長を引いた値」が等しければ条件を満たすということです。コードにすると以下の様な感じになります。

input()
A = [int(h) for h in input().split()]
iph_dic = {}
pair_of_spies = 0

for i, h in enumerate(A):
   if i - h in iph_dic:
       pair_of_spies += iph_dic[i - h]
   if i + h in iph_dic:
       iph_dic[i + h] += 1
   else:
       iph_dic[i + h] = 1

print(pair_of_spies)

こうやって、解答見て書いてみると簡単なんですが、問題からここに行き着くのはなかなか難しい。ただ、今までのE問題と比べると簡単な気がしますね。（まだ３回しかやってないけど）