公式

E - チーム編成 / Team Formation 解説 by admin

GPT 5.2 High

概要

\(N \le 30\) 人からちょうど \(K\) 人を選び、実力値の合計が \(M\) の倍数(\(\bmod M = 0\))になる選び方の数を数えます。\(N\) が小さめなので、半分全列挙(Meet-in-the-Middle)で高速に数え上げます。

考察

重要な気づき

選んだ集合の「合計が \(M\) の倍数」という条件は、合計の 剰余(\(\bmod M\) だけ見ればよい、という点が本質です。

また \(N \le 30\) なので、全ての組合せは \({30 \choose 15} \approx 1.55 \times 10^8\) 程度にもなり得て、素朴に「全ての \(K\) 人の組合せを列挙して和を計算」すると厳しいです(Python では特に TLE になりやすい)。

どう解決するか(半分全列挙)

配列を左右 2 つに分けます:

  • 左:\(n_1 = \lfloor N/2 \rfloor\)
  • 右:\(n_2 = N - n_1\)

それぞれについて「部分集合」を全列挙すると、個数は \(2^{15} = 32768\) 程度(最大でも)なので余裕があります。

ここで、 - 左から \(s\) 人選んだときの合計の剰余を \(r_L\) - 右から \(t\) 人選んだときの合計の剰余を \(r_R\) とすると、全体で \(K\) 人選ぶには \(s+t=K\) が必要で、さらに条件は [ (r_L + r_R) \bmod M = 0 ] です。つまり [ r_R \equiv -r_L \pmod M ] となる右側の剰余の個数を数えればよい、という「突き合わせ」問題になります。

具体例(イメージ)

例えば \(M=7\)、左で \(s=2\) 人選んだ部分集合の和の剰余が \(r_L=3\) なら、右で \(t=K-2\) 人選んだ部分集合の剰余は \((-3)\bmod 7 = 4\) である必要があります。
左側で「剰余 3 のものが \(c\) 個」、右側で「剰余 4 のものが \(d\) 個」なら、この組み合わせは \(c \times d\) 通りです。

アルゴリズム

  1. 配列 \(A\) を左右に分割する(left, right)。
  2. それぞれについて、全ての部分集合を列挙し、
    • 部分集合の要素数(何人選んだか)
    • 部分集合の和の剰余(\(\bmod M\)) を求めて、次の形で個数を数える: [ \text{cnt}[k][r] = \text{「その側から }k\text{ 人選び、和}\bmod M=r\text{ となる部分集合の個数」} ] 実装では dicts[k][r](辞書)として保持します(\(M\) が最大 \(10^9\) なので配列で \(0..M-1\) を持つのは不可能なため)。
  3. \(s\) を「左から選ぶ人数」として動かす:
    • \(s\) の取り得る範囲は、右側人数 \(n_2\) を使って [ \max(0, K-n_2) \le s \le \min(n_1, K) ]
    • \(t = K-s\) として、左 dL[s] と右 dR[t] を突き合わせる。
  4. 左の剰余 \(r\) に対して必要な右の剰余は \((-r)\bmod M\) なので、 [ \text{ans} += \text{dL}[s][r] \times \text{dR}[t][(-r)\bmod M] ] を足し込む(\(10^9+7\) で割った余りを管理)。
  5. 最終的な答えを出力する。

コード中の build の工夫

build では、全マスク(\(0..2^n-1\))に対し - lsb = mask & -mask(最下位ビット) - prev = mask ^ lsb(そのビットを外した集合) として、 [ \text{sums}[mask] = (\text{sums}[prev] + \text{追加された要素}) \bmod M ] [ \text{bits}[mask] = \text{bits}[prev] + 1 ] とすることで、各部分集合の「和 mod M」と「人数」を高速に求めています。

計算量

  • 時間計算量:
    部分集合列挙が左右それぞれ \(O(2^{n_1} + 2^{n_2})\)、突き合わせも「出現した剰余の数」だけなので概ね同程度です。
    よって全体で [ O(2^{N/2}) ] 程度(定数倍はありますが \(N \le 30\) なら十分間に合います)。
  • 空間計算量:
    左右の全部分集合情報を保持するため [ O(2^{N/2}) ] 程度です(sums, bits 配列と、サイズ別の辞書群)。

実装のポイント

  • \(M\) が非常に大きいので、剰余 \(0..M-1\) の配列は作れません。辞書(defaultdict(int))で「出現した剰余」だけ数えるのが重要です。

  • 条件 \(s+t=K\) を満たす \(s\) の範囲は必ず絞る(s_min, s_max)。不要なケースを回さないことで無駄が減ります。

  • 答えは組合せ数なので巨大になります。加算・乗算のたびに \(10^9+7\) で剰余を取ります。

  • 突き合わせでは、辞書サイズが小さい方を走査する(コード中の small, large の工夫)と高速になります。

    ソースコード

import sys
from collections import defaultdict

MOD = 10**9 + 7

def build(arr, M):
    n = len(arr)
    size = 1 << n
    sums = [0] * size
    bits = [0] * size

    for mask in range(1, size):
        lsb = mask & -mask
        i = lsb.bit_length() - 1
        prev = mask ^ lsb
        sums[mask] = (sums[prev] + arr[i]) % M
        bits[mask] = bits[prev] + 1

    dicts = [defaultdict(int) for _ in range(n + 1)]
    for mask in range(size):
        dicts[bits[mask]][sums[mask]] += 1
    return dicts

def main():
    input = sys.stdin.readline
    N, K, M = map(int, input().split())
    A = list(map(int, input().split()))

    n1 = N // 2
    left = A[:n1]
    right = A[n1:]

    dL = build(left, M)
    dR = build(right, M)

    ans = 0
    n2 = N - n1

    s_min = max(0, K - n2)
    s_max = min(n1, K)

    for s in range(s_min, s_max + 1):
        t = K - s
        d1 = dL[s]
        d2 = dR[t]

        if len(d1) <= len(d2):
            small, large = d1, d2
        else:
            small, large = d2, d1

        for r, c in small.items():
            ans = (ans + (c % MOD) * (large.get((-r) % M, 0) % MOD)) % MOD

    print(ans)

if __name__ == "__main__":
    main()

この解説は gpt-5.2-high によって生成されました。

投稿日時:
最終更新: