公式

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

GLM 5.2 (High)

概要

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

考察

全ての選び方を愚直に全探索しようとすると、最大で \(\binom{36}{18}\) 通りとなり、約90億通りになるため実行時間制限に間に合いません(TLEとなります)。

そこで、「半分全列挙」という手法を使います。配列を2つに分割し、それぞれで全探索を行い、その結果を効率よく組み合わせることで計算量を大幅に減らすことができます。

アルゴリズム

  1. 配列の分割: \(N\) 個のお菓子を前半 \(N_1 = \lfloor N/2 \rfloor\) 個と、後半 \(N_2 = N - N_1\) 個に分けます。
  2. それぞれの全列挙: 前半の配列と後半の配列それぞれについて、ビット全探索を用いて部分集合を列挙します。このとき、単に合計値を記録するだけでなく、「選んだお菓子の個数 \(i\)」ごとに合計値をまとめたリスト c1, c2 を作成します。
    • 例えば c1[i] には、前半から \(i\) 個選んだときの合計値がすべて格納されます。
  3. ソート: 二分探索を可能にするため、c2 のそれぞれのリストを昇順にソートします(c1 もソートしておきます)。
  4. 組み合わせの集計: 全体で \(K\) 個選ぶため、前半から \(i\) 個選ぶなら、後半からは \(j = K - i\) 個選ぶ必要があります。
    • \(i\) と \(j\) の範囲がそれぞれの配列長を超えないようにループを回します。
    • c1[i] に含まれる各合計値 \(a\) について、後半の合計値 \(b\) と合わせて \(X\) 以上になる必要があるため、\(b \geq X - a\) となります。
    • これは、ソート済みの c2[j] に対して二分探索を行い、\(X - a\) 以上となる要素の個数を求めることで高速に計算できます。その個数を答えに加算します。

計算量

  • 時間計算量: \(O(N \cdot 2^{N/2})\)
    • \(N \leq 36\) であるため、\(N/2 = 18\) となり、\(2^{18} \approx 260,000\) となります。全列挙と二分探索を合わせても十分に高速に実行可能です。
  • 空間計算量: \(O(2^{N/2})\)
    • 前半と後半の部分集合を記録するためのリストの要素数の合計がそれぞれ \(2^{N/2}\) 個になるためです。

実装のポイント

  • お菓子を「ちょうど \(K\) 個」選ぶ必要があるため、合計値のリストを「選んだ個数」ごとに分類して保持するのがポイントです。これにより、前半と後半の個数の和が \(K\) になるように簡単に組み合わせられます。

  • ビット全探索の際、lsb = m & -m と lsb.bit_length() - 1 を用いることで、立っているビットのインデックスを効率よく取得し、配列の添え字として使用しています。

  • Pythonの bisect モジュールを用いる際、bisect_left を使うことで「\(X - a\) 以上」という条件を満たす最小のインデックスを求め、そこから末尾までの要素数を計算しています。

    ソースコード

import sys
import bisect

def solve():
    input = sys.stdin.readline
    N, K, X = map(int, input().split())
    A = list(map(int, input().split()))
    
    N1 = N // 2
    N2 = N - N1
    
    L1 = A[:N1]
    L2 = A[N1:]
    
    c1 = [[] for _ in range(N1 + 1)]
    c1[0].append(0)
    for mask in range(1, 1 << N1):
        m = mask
        cnt = 0
        s = 0
        while m:
            lsb = m & -m
            idx = lsb.bit_length() - 1
            cnt += 1
            s += L1[idx]
            m ^= lsb
        c1[cnt].append(s)
        
    c2 = [[] for _ in range(N2 + 1)]
    c2[0].append(0)
    for mask in range(1, 1 << N2):
        m = mask
        cnt = 0
        s = 0
        while m:
            lsb = m & -m
            idx = lsb.bit_length() - 1
            cnt += 1
            s += L2[idx]
            m ^= lsb
        c2[cnt].append(s)
        
    for i in range(N1 + 1):
        c1[i].sort()
    for i in range(N2 + 1):
        c2[i].sort()
        
    ans = 0
    for i in range(max(0, K - N2), min(K, N1) + 1):
        j = K - i
        c2_j = c2[j]
        len_c2_j = len(c2_j)
        for a in c1[i]:
            target = X - a
            idx = bisect.bisect_left(c2_j, target)
            ans += len_c2_j - idx
            
    print(ans)

if __name__ == "__main__":
    solve()

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

投稿日時:
最終更新: