かしのブログ

競技プログラミングとか

AtCoder Beginner Contest 166 E - This Message Will Self-Destruct in 5s

AtCoder 変数分離

問題

$N$ 人がいて、 $i$ 番目の身長は $A_i$ である。「2人の番号の差の絶対値と、身長の和が等しい」ようなペアは何通りあるか。

考察

2人の番号を $i, j (i \lt j)$ とする。すると、ペアの制約は、 $j - i = A_i + A_j$ になる。番号 $i, j$ を分離すると、 $j - A_j = i + A_i$ となる。よって各番号 $i$ について、 $k \lt i \land i - A_i = k + A_k$ なる $k$ の数を数えればよい。

$M_{i, X} = |\{k | k \lt i \land X = k + A_k\}|$ とすると、解は $\sum_{i=1}^N M_{i, i+A_i}$ となる。 $i$ を小さい方から順に計算していくと、 $M_{i, X}$ の計算と合わせて $\mathcal{O}(N)$ で計算できる。実装では $M_{i, X}$ をmapで管理したので、計算量が $\mathcal{O}(N \log N)$ になった。

実装