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