BEN2のブログ

たまに書いています。

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

競技プログラミング

ABC 166 E - This Message Will Self-Destruct in 5s についてのメモを残します。

立式、変数の分離

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

参加者 $i$ と $j$ で取引できる条件は、 $|j - i| = A_{i} + A_{j}$ である。
ここで、 $i \lt j$ という条件をつけると、 $i + A_{i} = j - A_{j}$ と表せる。
$i \lt j$ かつ $i + A_{i} = j - A_{j}$ を満たす $(i, \, j )$ の組の総数を求めたい。
下図は、入力例 1 の $N = 6$ 、 $A = \{ 2, 3, 3, 1, 3, 1\}$ における $i + A_{i}$ と $j - A_{j}$ を表す。参加者の番号は $0, 1, \, \cdots , \, N-1$ としている。

f:id:BEN2suzuka:20200504153139p:plain

実装では、 $j$ を固定し、条件を満たす $i$ の個数を数える。
$j$ について、 $j - A_{j} = X$ とする。 $i \lt j$ の範囲で、 $i + A_{i}$ の値が $X$ となる $i$ の個数を数えるとよい。

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

int main() {
  int N; cin >> N;
  vector<int> A(N);
  for (int i = 0; i < N; i++) cin >> A.at(i);
  long long ans = 0;
  map<int, int> M;  // i + A[i] と その個数
  // j について全探索
  for (int i = 0; i < N; i++) {
    // j - A[j] = X とする
    int X = i - A.at(i);

    // i + A[i] = X となる i (i < j) の個数を加算
    if (M.count(X)) ans += M[X];

    // map に参加者 i の i + A[i] を追加
    if (M.count(i + A.at(i))) M[i + A.at(i)]++;
    else M[i + A.at(i)] = 1;
  }
  cout << ans << endl;
}