ABC 166 E - This Message Will Self-Destruct in 5s

問題

長さ $n$ の自然数列 $A$ に対して、$A _ i + A _ j = j - i$ となる $(i,j)$ の個数を求めよ。（$n \leq 2 \times 10 ^ 5$）

atcoder.jp

解説

$B _ i = A _ i + i, \ C _ j = j - A _ j$ とおいて、$B _ i = C _ j$ となる $(i,j)$ の個数を求めればよい。これは各 $i$ について $C _ j \leq B _ i$ なる $j$ の個数から $C _ j < B _ i$ なる $j$ の個数を引いたものを合計すればよいので、$C$ を昇順にソートした上で二分探索を使えば求められる。

実装

$C _ j \leq B _ i$ なる $j$ の個数と$C _ j < B _ i$ なる $j$ の個数はそれぞれ bisect.bisect_right と bisect.bisect_left で計算した。

import bisect
n = int(input())
*A, = map(int, input().split())
B = [A[i] + i for i in range(n)]
C = [i - A[i] for i in range(n)]
C.sort()
ans = 0
for b in B:
    ans += bisect.bisect_right(C, b) - bisect.bisect_left(C, b)
print(ans)

atcoder.jp