Official

E - お菓子の選抜 / Candy Selection Editorial by admin

GLM 5.2 (High)

概要

\(N\) 個のお菓子の中から \(K\) 個を選び、その満足度の合計が \(X\) 以上になるような選び方の総数を求める問題です。

考察

すべての選び方を全探索しようとすると、選び方の総数は最大で \(_{36}\mathrm{C}_{18} \approx 9 \times 10^9\) 通りにもなり、計算が終わりません(TLEになります)。 このような「要素数 \(N\) が 40 程度までの部分和問題」では、配列を2つに分けて考える半分全列挙が有効です。

配列を前半と後半の2つに分けます。それぞれの部分集合について「選んだ個数」と「満足度の合計」を全列挙します。前半で \(j\) 個選んだとき、後半では \(K - j\) 個選ぶ必要があります。前半で選んだ合計を \(s\) とすると、後半では \(X - s\) 以上の合計を持つ選び方を探すことで、条件を満たす組み合わせの総数を高速に求めることができます。

アルゴリズム

  1. 配列を分割: \(N\) 個の要素を前半 \(N_1 = N / 2\) 個と、後半 \(N_2 = N - N_1\) 個に分けます。
  2. 前半の全列挙: 前半の要素の選び方をビット全探索します。選んだ個数 \(j\) ごとに、その満足度の合計をリスト L_sums[j] に格納します。
  3. 後半の全列挙: 同様に後半の要素の選び方をビット全探索し、選んだ個数 \(j\) ごとにリスト R_sums[j] に満足度の合計を格納します。
  4. 後半リストのソート: 二分探索を可能にするため、各 R_sums[j] をソートしておきます。
  5. 条件を満たす組み合わせの計数:
    • 前半で \(j\) 個を選ぶ場合(\(j=0 \dots K\))を考えます。このとき後半では \(K-j\) 個を選ぶ必要があります。
    • L_sums[j] に格納された各合計値 \(s\) について、右側で必要な和は \(X - s\) 以上です。
    • R_sums[K - j] に対して二分探索を行い、\(X - s\) 以上の要素の個数を数え、答えに加算します。

計算量

  • 時間計算量: \(O(2^{N/2} \log(2^{N/2}))\) ビット全探索で \(O(2^{N/2})\)、ソートと二分探索で \(O(2^{N/2} \log(2^{N/2}))\) かかります。\(N=36\) のとき \(2^{18} \approx 2.6 \times 10^5\) なので、十分高速に処理できます。
  • 空間計算量: \(O(2^{N/2})\) 前半と後半の部分和を格納するリストのサイズに比例します。

実装のポイント

  • お菓子の満足度 \(A_i\) や、それらの合計値は \(10^9\) を超えることがあるため、C++などで実装する際は long long 型を使用してオーバーフローを防ぎます。

  • 二分探索を行う際は、ソート済みの配列に対して lower_bound を用いることで、「\(X - s\) 以上」という条件を満たす最初のイテレータを高速に取得できます。求める個数は、そのイテレータから配列の末尾までの距離(end() - it)で計算できます。

    ソースコード

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    
    int N, K;
    long long X;
    cin >> N >> K >> X;
    
    vector<long long> A(N);
    for (int i = 0; i < N; i++) {
        cin >> A[i];
    }
    
    int N1 = N / 2;
    int N2 = N - N1;
    
    vector<vector<long long>> L_sums(N1 + 1);
    vector<vector<long long>> R_sums(N2 + 1);
    
    for (int mask = 0; mask < (1 << N1); mask++) {
        long long sum = 0;
        int cnt = 0;
        for (int i = 0; i < N1; i++) {
            if (mask & (1 << i)) {
                sum += A[i];
                cnt++;
            }
        }
        L_sums[cnt].push_back(sum);
    }
    
    for (int mask = 0; mask < (1 << N2); mask++) {
        long long sum = 0;
        int cnt = 0;
        for (int i = 0; i < N2; i++) {
            if (mask & (1 << i)) {
                sum += A[N1 + i];
                cnt++;
            }
        }
        R_sums[cnt].push_back(sum);
    }
    
    for (int i = 0; i <= N2; i++) {
        sort(R_sums[i].begin(), R_sums[i].end());
    }
    
    long long ans = 0;
    for (int j = 0; j <= K; j++) {
        if (j > N1 || K - j > N2) continue;
        for (long long s : L_sums[j]) {
            long long target = X - s;
            auto it = lower_bound(R_sums[K - j].begin(), R_sums[K - j].end(), target);
            ans += R_sums[K - j].end() - it;
        }
    }
    
    cout << ans << "\n";
    
    return 0;
}

この解説は or-glm5.2-high によって生成されました。

posted:
last update: