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

問題情報

コンテスト：Codeforces #707
問題順：Div.1 A (Div.2 C)
問題へのリンク Div.1 Div.2

問題概要

数列$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";
}