Official

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

gemini-3.6-flash-high

概要

\(N\) 個のお菓子から \(K\) 個を選び、満足度の合計が \(X\) 以上となる組み合わせの数を求める問題です。\(N \le 36\) という制約に着目し、半分全列挙(Meet-in-the-middle)を用いることで高速に計算します。

考察

素朴なアプローチ(全探索)とその限界

まず考えるのは、\(N\) 個のお菓子から \(K\) 個を選ぶすべての組み合わせ(\({}_N \text{C}_K\) 通り)を試す方法です。 しかし、\(N = 36, K = 18\) の場合、組み合わせの数は \({}_{36} \text{C}_{18} = 90,751,353,000 \approx 9 \times 10^{10}\) 通りにもなり、実行時間制限(通常2秒程度)に到底間に合いません (TLE)。

動的計画法(DP)の検討

満足度の合計を状態に持つ動的計画法も考えられますが、お菓子の満足度 \(A_i\) が最大 \(10^9\) と非常に大きいため、配列のサイズが足りずメモリ・時間ともにオーバーしてしまいます。

半分全列挙による解決

ここで \(N \le 36\) という制約に注目します。 \(36\) 個の要素をそのまま全探索するのは不可能ですが、配列を前半 \(N_1 = \lfloor N/2 \rfloor\) 個(最大18個)と後半 \(N_2 = N - N_1\) 個(最大18個)の2つに半分ずつ分けるとどうでしょうか?

18個の要素からの部分集合の選び方は \(2^{18} = 262,144\) 通りしかありません。これならコンピュータにとって一瞬で計算できる大きさです。

したがって、以下の手順で解くことができます: 1. 前半のお菓子群から「\(k_1\) 個選んだときの満足度の和」をすべて列挙する。 2. 後半のお菓子群から「\(k_2\) 個選んだときの満足度の和」をすべて列挙する。 3. \(k_1 + k_2 = K\) となる組み合わせについて、(前半の和) + (後半の和) \(\ge X\) となる個数を効率よく数え上げる。

アルゴリズム

  1. お菓子の分割: お菓子の配列 \(A\) を前半 \(A_1\) (サイズ \(N_1\))と後半 \(A_2\) (サイズ \(N_2\))に分割します。

  2. 和の列挙: \(A_1, A_2\) それぞれについて、選んだ個数 \(c\) ごとの満足度の合計リストを作成します。

    • sums1[c] : 前半から \(c\) 個選んだときの和のリスト
    • sums2[c] : 後半から \(c\) 個選んだときの和のリスト
  3. ソート: 効率よくカウントするために、各 sums1[c] および sums2[c] を昇順にソートします。

  4. しゃくとり法(Two Pointers)によるカウント: 前半から選ぶ個数 \(k_1\) を \(0\) から \(N_1\) まで試します。このとき、後半から選ぶ個数は \(k_2 = K - k_1\) と決まります(\(0 \le k_2 \le N_2\) の場合のみ処理)。

s1(sums1[k1] の要素)と s2(sums2[k2] の要素)の和が \(X\) 以上になるペアの数を数えます。 s1 を小さい順に見るとき、s1 + s2 >= X を満たす最小の s2 の位置(インデックス \(j\))は、s1 が大きくなるにつれて左(小さい方)へ移動します。この性質を利用してポインタ \(j\) を右から左に動かす「しゃくとり法」を使うことで、各 \(k_1\) に対するカウントを高速に行うことができます。

計算量

  • 時間計算量: \(O(N \cdot 2^{N/2})\)

    • 前半・後半の列挙に \(O(2^{N/2})\) かかります。
    • ソートおよびしゃくとり法による比較に \(O(N_1 \cdot 2^{N/2}) = O(N \cdot 2^{N/2})\) かかります。
    • \(N = 36\) のとき \(2^{18} \approx 2.6 \times 10^5\) であり、\(N \cdot 2^{N/2} \approx 9.4 \times 10^6\) 回程度の計算となるため、余裕で実行時間制限内に収まります。
  • 空間計算量: \(O(2^{N/2})\)

    • 前半と後半の全列挙結果を保持するために \(2^{N_1} + 2^{N_2} \le 2 \times 2^{18} \approx 5.2 \times 10^5\) 個の要素をメモリに保持します。

実装のポイント

  • 個数ごとの管理: 単に満足度の和を列挙するだけでなく、「選んだお菓子の個数 \(c\)」ごとに分けて管理することが重要です。これにより、合計選択数がちょうど \(K\) 個になる組み合わせのみを適切に判定できます。

  • しゃくとり法の向き: s1 を昇順(小さい順)に試す場合、s1 + s2 >= X を満たす境界 \(j\) は s1 の増加に伴って小さくなります。そのため、ポインタ \(j\) を配列の末尾からスタートさせ、条件を満たす限り左へ減らしていくことで、\(O(\text{len}(s1) + \text{len}(s2))\) でカウントできます。

    ソースコード

import sys


def solve():
    input_data = sys.stdin.read().split()
    if not input_data:
        return

    N = int(input_data[0])
    K = int(input_data[1])
    X = int(input_data[2])

    A = [int(x) for x in input_data[3:]]

    N1 = N // 2
    N2 = N - N1

    A1 = A[:N1]
    A2 = A[N1:]

    def get_sums(arr):
        n = len(arr)
        res = [[] for _ in range(n + 1)]
        res[0].append(0)
        for x in arr:
            for c in range(n - 1, -1, -1):
                if res[c]:
                    res[c + 1].extend([s + x for s in res[c]])
        return res

    sums1 = get_sums(A1)
    sums2 = get_sums(A2)

    for c in range(N1 + 1):
        sums1[c].sort()
    for c in range(N2 + 1):
        sums2[c].sort()

    ans = 0

    for k1 in range(N1 + 1):
        k2 = K - k1
        if 0 <= k2 <= N2:
            s1_list = sums1[k1]
            s2_list = sums2[k2]
            len2 = len(s2_list)

            j = len2
            for s1 in s1_list:
                while j > 0 and s1 + s2_list[j - 1] >= X:
                    j -= 1
                ans += len2 - j

    print(ans)


if __name__ == "__main__":
    solve()

この解説は gemini-3.6-flash-high によって生成されました。

posted:
last update: