Official

E - チーム編成 / Team Formation Editorial by admin Editorial for printing

Claude 4.5 Opus

概要

\(N\) 人の生徒から \(K\) 人を選び、実力値の合計が \(M\) の倍数となる選び方の数を求める問題です。半分全列挙(Meet in the Middle)を使って効率的に解きます。

考察

素朴なアプローチの問題点

\(N\) 人から \(K\) 人を選ぶ全ての組み合わせを調べる方法を考えます。組み合わせの数は \(\binom{N}{K}\) 通りで、最大で \(\binom{30}{15} \approx 1.5 \times 10^8\) 通りにもなります。さらに各組み合わせで合計を計算すると、全体で \(O(2^N \cdot N)\) となり、TLE(時間制限超過)になる可能性が高いです。

重要な気づき

\(N \leq 30\) という制約は「半分全列挙」を示唆しています。\(2^{30} \approx 10^9\) は厳しいですが、\(2^{15} \approx 3 \times 10^4\) なら十分高速に処理できます。

解決方針

生徒を前半グループと後半グループの2つに分けます。それぞれのグループで全ての部分集合を列挙し、「何人選んだか」と「合計を \(M\) で割った余り」を記録しておきます。その後、両グループの結果を組み合わせて、合計 \(K\) 人で合計が \(M\) の倍数になる選び方を数えます。

アルゴリズム

  1. グループ分割: \(N\) 人を前半(\(\lfloor N/2 \rfloor\) 人)と後半(残り)に分ける

  2. 部分集合の列挙: 各グループについて、全ての部分集合を列挙し、以下を記録する

    • 選んだ人数 \(k\)
    • 実力値の合計を \(M\) で割った余り \(r\)
    • そのような部分集合の個数
  3. 結果の統合: 前半から \(k_1\) 人、後半から \(k_2 = K - k_1\) 人選ぶ組み合わせを考える

    • 前半の余りが \(r_1\) のとき、後半の余りが \(r_2 = (M - r_1) \mod M\) であれば合計が \(M\) の倍数
    • 条件を満たす組み合わせ数を掛け合わせて足し上げる

具体例

\(N=4, K=2, M=3, A=[1,2,3,4]\) の場合: - 前半 \([1,2]\)、後半 \([3,4]\) に分割 - 前半から1人選ぶ: \((1, 余り1), (1, 余り2)\) が各1通り - 後半から1人選ぶ: \((1, 余り0), (1, 余り1)\) が各1通り - \(k_1=1, k_2=1\) で余りの和が \(M\) の倍数になる組み合わせを数える - 前半余り1 × 後半余り2 → 該当なし - 前半余り2 × 後半余り1 → 該当なし - など、全パターンをチェック

計算量

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

    • 各グループで \(2^{N/2}\) 個の部分集合を列挙
    • 統合処理は部分集合の数に依存するが、最大でも \(O(2^{N/2} \cdot 2^{N/2} / M)\) 程度
    • \(N=30\) のとき約 \(2 \times 2^{15} \approx 6 \times 10^4\) 回の列挙
  • 空間計算量: \(O(2^{N/2})\)

    • 各グループの部分集合情報を辞書で保持

実装のポイント

  • ビットマスクによる列挙: mask の各ビットが要素の選択/非選択を表す

  • 辞書の活用: result[k][r] で「\(k\) 人選び、余り \(r\)」となる通り数を管理

  • 余りの補数計算: 合計が \(M\) の倍数になる条件は \(r_1 + r_2 \equiv 0 \pmod{M}\)、すなわち \(r_2 = (M - r_1) \mod M\)

  • \(10^9+7\) での剰余: 答えが大きくなるため、掛け算・足し算のたびに剰余を取る

    ソースコード

import sys
from collections import defaultdict

def solve():
    input_data = sys.stdin.read().split()
    N = int(input_data[0])
    K = int(input_data[1])
    M = int(input_data[2])
    A = [int(input_data[3 + i]) for i in range(N)]
    
    MOD = 10**9 + 7
    
    # Meet in the middle approach
    # Split N people into two halves
    half = N // 2
    left = A[:half]
    right = A[half:]
    
    # For left half: enumerate all subsets
    # Store (count of people, sum mod M) -> number of ways
    def enumerate_subsets(arr):
        n = len(arr)
        # result[k][r] = number of subsets with k people and sum % M == r
        result = defaultdict(lambda: defaultdict(int))
        for mask in range(1 << n):
            cnt = 0
            total = 0
            for i in range(n):
                if mask & (1 << i):
                    cnt += 1
                    total += arr[i]
            result[cnt][total % M] += 1
        return result
    
    left_subsets = enumerate_subsets(left)
    right_subsets = enumerate_subsets(right)
    
    ans = 0
    
    # Combine: choose k1 from left, k2 from right where k1 + k2 = K
    # sum1 + sum2 ≡ 0 (mod M) => sum2 ≡ -sum1 (mod M)
    
    for k1 in range(K + 1):
        k2 = K - k1
        if k1 > len(left) or k2 > len(right):
            continue
        if k2 < 0:
            continue
        
        left_dict = left_subsets[k1]
        right_dict = right_subsets[k2]
        
        for r1, cnt1 in left_dict.items():
            # Need r2 such that (r1 + r2) % M == 0
            r2 = (M - r1) % M
            if r2 in right_dict:
                ans = (ans + cnt1 * right_dict[r2]) % MOD
    
    print(ans)

solve()

この解説は claude4.5opus によって生成されました。

posted:
last update: