【典型90問】051 - Typical Shop（★5）

atcoder.jp

制約を見て半分全列挙とかいうやつじゃないか、と気づいた。
勉強したことなかったので記事にします。

半分全列挙について

簡単にいうとbit全探索もしくは何重かfor文全探索を、二分探索などで高速化しようという考え方のこと。

下の記事の解説がとてもわかりやすかった。
algo-logic.info

要素を半分に分割してそれぞれ全列挙し、その後片方を固定して、二分探索などでもう片方から答えを探し出す。

実装

$ {}_n C_k $ の探索をどうやって実装しようかと思ったが、bit全探索で選んだ個数ごとに二次元配列に突っ込めば上手くできた。
Nが奇数の可能性があるので、bit全探索は2回に分けて書いている。
工夫すれば一回でもいけそう。
二分探索はupper_boundを使って、n以下の要素が何個あるか数えるやつをしている。

void solve(ll N, ll K, ll P, std::vector<long long> A) {
    // 半分に分割
    vector<ll> A1, A2;
    REP (i, N) {
        if (i < N/2) A1.push_back(A[i]);
        else A2.push_back(A[i]);
    }

    // 選んだ個数ごとに格納
    vector<vector<ll>> sm1(K+1), sm2(K+1);
    
    // 前半列挙
    for (int bit = 0; bit < (1<<(N/2)); bit++) {
        ll cur = 0;
        int cnt = 0;
        for (int i = 0; i < N/2; i++) {
            if (bit & (1<<i)) {
                cnt++;
                cur += A1[i];
            }
        }

        if (cnt <= K) sm1[cnt].push_back(cur);
    }

    // 後半列挙
    for (int bit = 0; bit < (1<<myceil(N, 2LL)); bit++) {
        ll cur = 0;
        int cnt = 0;
        for (int i = 0; i < myceil(N, 2LL); i++) {
            if (bit & (1<<i)) {
                cnt++;
                cur += A2[i];
            }
        }

        if (cnt <= K) sm2[cnt].push_back(cur);
    }

    // 二分探索
    ll ans = 0;
    for (int i = 0; i <= K; i++) {
        ALL(sort, sm2[K-i]);
        REP (j, sm1[i].size()) {
            ans += upper_bound(all(sm2[K-i]), P-sm1[i][j]) - sm2[K-i].begin();
        }
    }

    c(ans)
}