C - 部活動の選択 / Choosing a Club Activity 解説 by admin
gemini-3-flash-thinking概要
\(N\) 個のアイテム(部活動)からいくつかを選び、その合計値を \(K\) の倍数にする組み合わせの数を求める問題です。 制約として \(N\) が最大 \(10^5\) と大きい一方で、\(K\) は \(100\) 以下と非常に小さいことが特徴です。
考察
1. 基本的な動的計画法 (DP)
まず思いつくのは、「\(dp[i][j] = i\) 番目までの部活動を考えて、活動ポイントの合計を \(K\) で割った余りが \(j\) になる組み合わせの数」という DP です。 しかし、このままでは計算量が \(O(N \times K)\) となり、今回の制約(\(N=10^5, K=100\))では \(10^7\) 程度の計算回数になります。Python の場合、単純なループでは実行時間制限に間に合わない可能性があるため、より効率的なアプローチを考えます。
2. 生成関数と多項式の利用
各部活動を選ぶ・選ばないという選択は、多項式 \((1 + x^{A_i})\) で表現できます。 すべての部活動の組み合わせは、以下の多項式の積として表せます: $\(P(x) = \prod_{i=1}^N (1 + x^{A_i})\)\( この多項式を展開したとき、\)x^m\( の係数は「合計ポイントが \)m\( になる組み合わせの数」です。 私たちが知りたいのは「合計が \)K\( の倍数」になる数なので、多項式の次数を \)K\( で割った余りで管理すれば十分です。つまり、**多項式の積を \)x^K - 1 \equiv 0\((すなわち \)x^K \equiv 1$)という関係式の下で計算する**ことになります。これは「巡回畳み込み」と呼ばれる操作と同じです。
3. 計算の高速化
同じ余りを持つ部活動をまとめることで、計算を大幅に短縮できます。 - 活動ポイントを \(K\) で割った余りが \(r\) である部活動が \(c\) 個ある場合、その部分の積は \((1 + x^r)^c\) となります。 - \((1 + x^r)^c \pmod{x^K - 1}\) を計算するには、まず \((1 + x)^c \pmod{x^K - 1}\) を繰り返し二乗法で求め、その後に \(x\) を \(x^r\) に置き換えることで高速に計算可能です。
アルゴリズム
- 余りごとのカウント: 各 \(A_i\) を \(K\) で割った余り \(r\) を数えます(
counts[r])。 - 特殊なケース (\(r=0\)): 余りが \(0\) の部活動が \(c_0\) 個ある場合、これらはどのように選んでも合計の余りに影響しません。各部活動について「入る・入らない」の 2 通りがあるため、全体に \(2^{c_0}\) を掛け合わせることで処理できます。
- 繰り返し二乗法による多項式累乗:
- \(r \in \{1, \dots, K-1\}\) について、出現する回数 \(c = counts[r]\) の種類ごとに \((1 + x)^c \pmod{x^K - 1}\) を計算します。
- 二分累乗法を用いることで、多項式の積を \(O(K^2 \log N)\) で計算できます。
- 多項式の合成:
- 求めた \((1 + x)^c\) の結果から、各 \(r\) に対して \((1 + x^r)^c\) を構成します。
- これらをすべて掛け合わせ、最終的な多項式 \(dp\) を得ます。
- 答えの出力:
- \(dp[0]\) が「合計が \(K\) の倍数(余り 0)」になる組み合わせの数です。
- ここには「何も選ばない(空集合)」の 1 通りが含まれているため、最後に \(1\) を引いて \(10^9+7\) で割った余りを出力します。
計算量
- 時間計算量: \(O(N + K^3 \log N)\)
- 最初のカウントに \(O(N)\)。
- 多項式の累乗(二分累乗法)に \(O(K^2 \log N)\)。これを最大 \(K\) 種類行うため \(O(K^3 \log N)\)。
- 多項式の乗算(巡回畳み込み)に \(O(K^3)\)。
- \(K=100\) なので \(K^3 = 10^6\) 程度となり、十分に高速です。
- 空間計算量: \(O(K^2)\)
- 累乗結果のキャッシュや DP テーブルの保持に必要なメモリです。
実装のポイント
多項式乗算の最適化: Python では内側のループで毎回
% MODを取ると遅くなるため、合計を計算してから最後に余りを取る工夫をしています。\(K=1\) の処理: \(K=1\) のときはすべての組み合わせが条件を満たします。\(2^N - 1\) を直接計算して出力します。
負の余りの処理:
dp[0] - 1が負になる可能性があるため、(dp[0] - 1) % MODとすることで正しく余りを計算します。ソースコード
import sys
# The problem asks for the number of non-empty subsets of clubs such that the sum of their activity points is a multiple of K.
# We use dynamic programming with generating functions.
# For each remainder r (mod K), let c be the number of clubs with activity point A_i % K = r.
# The generating function for these clubs is (1 + x^r)^c mod (x^K - 1).
# We can compute this as W_c(x^r) mod (x^K - 1), where W_c(x) = (1 + x)^c mod (x^K - 1).
# Since K is small (K <= 100), we can perform polynomial multiplication and exponentiation efficiently.
def solve():
# Fast I/O: Read all input at once
data = sys.stdin.read().split()
if not data:
return
N = int(data[0])
K = int(data[1])
A = list(map(int, data[2:]))
MOD = 10**9 + 7
if K == 1:
# If K=1, any non-empty subset of N clubs works.
# Total number of non-empty subsets is 2^N - 1.
print((pow(2, N, MOD) - 1) % MOD)
return
# Count frequencies of each remainder modulo K.
counts = [0] * K
for a in A:
counts[a % K] += 1
# Polynomial multiplication modulo (x^K - 1).
# This is equivalent to a cyclic convolution of length K.
def poly_mul(P, Q, K, MOD):
res = [0] * K
for i in range(K):
pi = P[i]
if pi == 0:
continue
# Summing products without modulo to speed up the inner loop.
# Max value of res[k] is approx K * MOD^2, which is ~10^20.
# Python handles arbitrary precision integers automatically.
for j in range(K - i):
res[i + j] += pi * Q[j]
for j in range(i):
res[j] += pi * Q[K - i + j]
# Apply modulo once for each coefficient at the end of the multiplication.
return [v % MOD for v in res]
# Polynomial exponentiation modulo (x^K - 1).
def poly_pow(P, n, K, MOD):
res = [0] * K
res[0] = 1
base = P
while n > 0:
if n % 2 == 1:
res = poly_mul(res, base, K, MOD)
base = poly_mul(base, base, K, MOD)
n //= 2
return res
# dp[j] is the number of ways to choose a subset with sum congruent to j mod K.
# Initially, only the empty set exists with sum 0.
dp = [0] * K
dp[0] = 1
# Precompute W_c(x) = (1 + x)^c mod (x^K - 1) for distinct counts c.
# This optimization reduces the number of poly_pow calls to at most K-1.
distinct_counts = set(counts[1:])
if 0 in distinct_counts:
distinct_counts.remove(0)
W_cache = {}
for c in distinct_counts:
P_init = [0] * K
P_init[0] = 1
P_init[1] = 1
W_cache[c] = poly_pow(P_init, c, K, MOD)
# Handle clubs with activity points divisible by K separately.
# Each of the counts[0] clubs can either be included or not.
if counts[0] > 0:
factor = pow(2, counts[0], MOD)
for j in range(K):
dp[j] = (dp[j] * factor) % MOD
# Handle other remainders r = 1, 2, ..., K-1.
for r in range(1, K):
c = counts[r]
if c == 0:
continue
W_c = W_cache[c]
# V_r(x) = (1 + x^r)^c mod (x^K - 1).
# This is equivalent to replacing x with x^r in W_c(x) mod (x^K - 1).
# Since (1 + x^r)^c = sum_{k=0}^c binom(c, k) x^{rk},
# and W_c(x) = sum_{j=0}^{K-1} w_j x^j where w_j = sum_{k=j mod K} binom(c, k),
# it follows that (1 + x^r)^c = sum_{j=0}^{K-1} w_j x^{rj mod K}.
V_r = [0] * K
for j in range(K):
V_r[(j * r) % K] += W_c[j]
# Modulo V_r coefficients as multiple j's might map to the same index.
V_r = [v % MOD for v in V_r]
# Combine the current possibilities with clubs of remainder r.
dp = poly_mul(dp, V_r, K, MOD)
# The result is dp[0] (total ways for sum = 0 mod K) minus 1 (to exclude the empty set).
# Python's modulo operator handles negative results correctly.
print((dp[0] - 1) % MOD)
if __name__ == '__main__':
solve()
この解説は gemini-3-flash-thinking によって生成されました。
投稿日時:
最終更新: