公式

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\) に置き換えることで高速に計算可能です。

アルゴリズム

  1. 余りごとのカウント: 各 \(A_i\)\(K\) で割った余り \(r\) を数えます(counts[r])。
  2. 特殊なケース (\(r=0\)): 余りが \(0\) の部活動が \(c_0\) 個ある場合、これらはどのように選んでも合計の余りに影響しません。各部活動について「入る・入らない」の 2 通りがあるため、全体に \(2^{c_0}\) を掛け合わせることで処理できます。
  3. 繰り返し二乗法による多項式累乗:
    • \(r \in \{1, \dots, K-1\}\) について、出現する回数 \(c = counts[r]\) の種類ごとに \((1 + x)^c \pmod{x^K - 1}\) を計算します。
    • 二分累乗法を用いることで、多項式の積を \(O(K^2 \log N)\) で計算できます。
  4. 多項式の合成:
    • 求めた \((1 + x)^c\) の結果から、各 \(r\) に対して \((1 + x^r)^c\) を構成します。
    • これらをすべて掛け合わせ、最終的な多項式 \(dp\) を得ます。
  5. 答えの出力:
    • \(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 によって生成されました。

投稿日時:
最終更新: