hotaruの蛍雪日記

競プロとゲームをしています

C. Going Home - Codeforces #707 Div.1 A (Div.2 C)

問題情報

問題概要

数列$a_1, a_2, \ldots, a_n$がある。 $a_x + a_y = a_z + a_w$となる一意な添字$x,y,z,w$は存在するか。

制約:$4 \leq n \leq 2 \times 10 ^ 5$、$1 \leq a_i \leq 2.5 \times 10 ^ 6$。

コード

折りたたみを開いてコードを見る

#include <iostream>
#include <vector>
using namespace std;

constexpr int MAX_A = 2.5 * 1000000;

// 2組の数値はすべて一意かどうか
// FIXME: 同じ組に同じ値があるとき正しい答えを返さない
constexpr bool isIndicesUnique(pair<int, int> a, pair<int, int> b) {
    return a.first != b.first && a.first != b.second && a.second != b.first && a.second != b.second;
}

int main() {
    // 入力受け取り
    int n;
    cin >> n;
    vector<int> a(n);
    for (int i = 0; i < n; i++) cin >> a[i];

    // 探索結果保存配列 buf[sum]: 和がsumのときの添字の組 -1はまだ保存されていないことを表す
    vector<pair<int, int>> buf(2 * MAX_A + 1, {-1, -1});

    // 全探索
    // i: 1つめの添字
    for (int i = 0; i < n; i++) {
        // j: 2つめの添字
        for (int j = i + 1; j < n; j++) {
            const int sum = a[i] + a[j];
            const pair<int, int> indices = {i, j};

            if (buf[sum].first == -1) {
                // まだこの和の添字は保存されていなかった→保存する
                buf[sum] = indices;
            } else if (isIndicesUnique(buf[sum], indices)) {
                // この和の添字は保存されていて、今のと以前のは一意だった→2つの添字の組が答え
                // 答え出力してプログラム終了
                cout << "YES\n";
                cout << indices.first + 1 << " " << indices.second + 1 << " " << buf[sum].first + 1 << " "
                     << buf[sum].second + 1 << "\n";
                return 0;
            }
        }
    }

    // 答えが見つからなかった プログラム終了
    cout << "NO\n";
}

解き方

全探索・鳩の巣原理

  1. $x, y$の組み合わせを全探索する。そのとき$a_x + a_y = sum$としよう
  2. 長さ$5 \times 10 ^ 6$の配列bufを用意する
  3. buf[sum]に$(x, y)$を記録する
  4. もしbuf[sum]にもう記録があって、探索中の$(x, y)$と記録の$(x, y)$のすべての添字が一意ならば答えが見つかった
  5. もし答えが見つからずに全探索が終了したならば答えが見つからなかった
  6. 全探索は高々$3 \times 5 \times 10 ^ 6 + 1$回で終わると思う
  7. $sum$の範囲は$2 \leq sum \leq 5 \times 10 ^ 6$である
  8. $sum$となる$x, y$の組が4組あれば必ず答えが存在する
  9. すべての$sum$で$x, y$の組を3組探索しても探索回数は$3 \times 5 \times 10 ^ 6$で次の探索結果は必ず答え

感想

  • コンテスト中は回答できなかった
  • これ鳩の巣原理じゃん
  • 探索回数の上限もっと減るはずだがよくわからんな
  • Editorialは $\mathrm{O} (\min(n ^ 2, \max(sum)))$ と言ってた

類題がある。