【競プロ】ABC200 D問題

競技プログラミング
Page content

まえがき

今後AtCoderにてABCのD問題やARCのB問題が解けなかった時に記事を書くことにしました。

今回はABC200の「D - Happy Birthday! 2」をコンテスト中に解けなかったため記事を書いています。
公式解説をもとに、コードを加えつつ再解説していきます。

問題の概要

N個の自然数からなる数列Aが与えられます。
この数列の空でない部分列について、総和の200の剰余が等しくなる部分列のペアを見つけるというのが今回のタスクです。
Nの範囲は$[2, 200]$、Aに含まれる整数の範囲は$[1, 10^9]$です。

例を挙げると入力が以下とすると、

3
4 8 212

出力は以下のようになります。

Yes
2 1 2
1 3

出力の1行目はペアが存在する場合"Yes"に、存在しない場合は"No"になります。
2行目と3行目では初めに部分列の長さを出力し、続けて部分列の要素の数列Aでの番号を出力します。
上の例では、数列Aが$(4, 8, 212)$で、部分列のペアは$(4, 8)$と$(212)$です。

今回の敗因

今回の敗因は全探索時の最悪計算量が$2^N$になると考えたことです。
$2^N$ではNが最大のとき$2^{200}$となり、単純な全探索では間に合わないと思いました。

後述しますが、実際には最悪計算量は$min(2^N, 201)$となり、全探索でも十分に間に合います。
部分列を201通り以上探索できる場合は必ず解となるペアが見つかります。

公式の解説の概要

201通り以上探索できた場合に解が見つかる理由として、解説では「鳩の巣原理」で説明されています。
「鳩の巣原理」は、$n$個のものを$m$個の箱に入れるとき、$n>m$であれば、2個以上のものが入っている箱が存在する、という原理です。

今回のタスクでは部分列を探索し、総和の200の剰余が等しいペアを求めます。
総和の200の剰余の範囲は$[0, 199]$であり、全ての部分列はこの範囲内のいずれかの整数に対応します。
「鳩の巣原理」で言えば$m=200$のため$n>200$であれば今回のタスクの解となるペアが見つかります。
$n$は部分列の探索数であり、部分列を201通り以上探索すると必ず解のペアが見つかることになります。

ACできたコード

#include <bits/stdc++.h>

#define rep(i, end) for (auto i = 0; i < end; ++i)

using namespace std;

vector<int> bits_to_indexes(const bitset<8>& bits) {
    vector<int> indexes;
    rep(i, 8) if (bits[i]) indexes.emplace_back(i);

    return indexes;
}

int sum_by_indexes(const vector<int>& indexes, const vector<int>& A) {
    int sum_a = 0;
    for (int i : indexes) sum_a += A[i];

    return sum_a;
}

void output_indexes(const vector<int>& indexes) {
    cout << indexes.size();
    for (int i : indexes) cout << ' ' << i + 1;
    cout << endl;
}

int main() {
    // Input
    int N;
    cin >> N;
    vector<int> A(N);
    rep(i, N) cin >> A[i];

    // Process
    rep(i, N) A[i] %= 200;

    vector<int> bits_memory(200, 0);  // 探索済bit列記憶用リスト
    bool judge = false;
    int b, c;
    rep(i, pow(2, N)) {  // 鳩の巣原理より、iが201を超える事はない
        bitset<8> bits(i);

        int sum_a = sum_by_indexes(bits_to_indexes(bits), A) % 200;

        if (bits_memory[sum_a] == 0) {
            // 総和の200の剰余が等しい部分列が未発見の場合、記憶用リストにbit列を表す数を格納
            bits_memory[sum_a] = i;
        } else {
            // 総和の200の剰余が等しい部分列が発見済の場合、bに探索済bit列を表す数、cに現在のbit列を表す数を格納してループを終える
            b = bits_memory[sum_a];
            c = i;
            judge = true;
            break;
        }
    }

    // Output
    if (judge) {
        cout << "Yes" << endl;

        vector<int> B = bits_to_indexes((bitset<8>)b);
        output_indexes(B);

        vector<int> C = bits_to_indexes((bitset<8>)c);
        output_indexes(C);
    } else {
        cout << "No" << endl;
    }

    return 0;
}