Official

E - 商品の逆元ポイント / Modular Inverse Points of Products Editorial by admin

Claude 4.6 Opus (Thinking)

概要

\(N\) 個の商品から \(K\) 個を選ぶ全ての組み合わせについて、選んだ商品の価格の積の「\(M\) を法とした乗法逆元」の総和を求める問題です。これは、各商品の価格の逆元に対する 第 \(K\) 次基本対称式 を計算する問題に帰着できます。

考察

重要な気づき 1: \(M\) の倍数の扱い

\(B_i = A_i \mod M\) とします。\(B_i = 0\) の商品が選ばれた場合、その組み合わせの積 \(P\) は必ず \(M\) の倍数になるため、ポイントは \(0\) です。よって、\(B_i \neq 0\) の商品だけを考えればよいです。

重要な気づき 2: 逆元の積への分解

\(M\) が素数なので、\(B_i \neq 0\) の商品だけからなる部分集合について:

\[P^{-1} \equiv \prod_{i \in S} B_i^{-1} \pmod{M}\]

が成り立ちます。つまり「積の逆元」は「逆元の積」に等しいです。

重要な気づき 3: 基本対称式への帰着

求めたい値は:

\[\sum_{|S|=K} \prod_{i \in S} B_i^{-1} \pmod{M}\]

これは \(c_i = B_i^{-1} \mod M\) として、\(c_1, c_2, \ldots, c_n\) の 第 \(K\) 次基本対称式 \(e_K(c_1, \ldots, c_n)\) に他なりません。

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

全ての \(\binom{N}{K}\) 通りを列挙すると、\(N=5000, K=2500\) のとき天文学的な数になり、到底間に合いません。

アルゴリズム

基本対称式はDPで効率的に計算できます。

DP定義: \(dp[j]\) = これまでに見た要素から \(j\) 個選んだときの、選んだ要素の積の総和(\(\mod M\))

初期状態: \(dp[0] = 1\), それ以外は \(0\)

遷移: 新しい要素 \(x\) を追加するとき(\(j\) を大きい方から更新): $\(dp[j] \leftarrow dp[j] + dp[j-1] \cdot x \pmod{M}\)$

\(j\) を大きい方から更新するのは、同じ要素を複数回使わないためです(0-1ナップサック型のテクニック)。

具体例: \(c = [2, 3, 5]\), \(K=2\) の場合 - 初期: \(dp = [1, 0, 0]\) - \(x=2\) 追加後: \(dp = [1, 2, 0]\) - \(x=3\) 追加後: \(dp = [1, 5, 6]\) - \(x=5\) 追加後: \(dp = [1, 10, 31]\)

\(dp[2] = 31 = 2 \cdot 3 + 2 \cdot 5 + 3 \cdot 5\) となり、正しく基本対称式が計算されています。

計算量

  • 時間計算量: \(O(N \cdot K + N \log M)\)
    • DPの更新: \(O(N \cdot K)\)(最大 \(5000 \times 5000 = 25 \times 10^6\))
    • 各要素の逆元計算: \(O(N \log M)\)(フェルマーの小定理による \(\text{pow}(b, M-2, M)\))
  • 空間計算量: \(O(N + K)\)

実装のポイント

  • 逆元の計算: \(M\) が素数なので、フェルマーの小定理より \(b^{-1} \equiv b^{M-2} \pmod{M}\) で計算できます。Pythonの pow(b, M-2, M) は高速です。

  • DPの更新順序: \(j\) を \(K\) から \(1\) へ降順に更新することで、同一要素の重複使用を防ぎます。これは0-1ナップサック問題と同じテクニックです。

  • 非零要素が \(K\) 未満の場合: \(B_i \neq 0\) の商品が \(K\) 個未満なら、\(K\) 個全てを非零で選ぶことが不可能なので答えは \(0\) です。

    ソースコード

import sys
from collections import Counter

def solve():
    input_data = sys.stdin.buffer.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)]
    
    # We need to compute sum over all K-subsets S of inv(prod(A_i for i in S)) mod M
    # where inv is modular inverse mod M, and if prod is 0 mod M, contribution is 0.
    
    # Let B_i = A_i mod M. Some B_i might be 0 (when A_i is divisible by M).
    # If B_i = 0, then any subset containing item i has product divisible by M, contributing 0.
    # So only subsets consisting entirely of items with B_i != 0 matter.
    
    # For items with B_i != 0, inv(prod) = prod(inv(B_i)).
    # So we need: sum over all K-subsets of non-zero items of prod(inv(B_i) for i in subset).
    
    # This is the K-th elementary symmetric polynomial of {inv(B_i)} for non-zero items.
    
    B = [a % M for a in A]
    
    # Separate zero and non-zero
    non_zero = [b for b in B if b != 0]
    n = len(non_zero)  # number of non-zero items
    
    # If n < K, all subsets either have fewer than K non-zero items or include zero items
    # Either way, answer is 0
    if n < K:
        print(0)
        return
    
    # Compute inv(b) for each non-zero b
    # inv(b) = pow(b, M-2, M)
    inv_vals = [pow(b, M - 2, M) for b in non_zero]
    
    # We need e_K(inv_vals) = sum of products of all K-subsets of inv_vals
    # This is the K-th elementary symmetric polynomial.
    
    # Standard DP: dp[j] = e_j after processing some elements
    # dp[0] = 1, and for each new element x, update dp[j] = dp[j] + dp[j-1]*x for j from K down to 1
    
    # n can be up to 5000, K up to 5000, so O(n*K) which is up to 25*10^6 - should be okay
    
    dp = [0] * (K + 1)
    dp[0] = 1
    
    for x in inv_vals:
        # Update from high to low to avoid using same element twice
        for j in range(min(K, dp.__len__() - 1), 0, -1):
            dp[j] = (dp[j] + dp[j - 1] * x) % M
    
    print(dp[K] % M)

solve()

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

posted:
last update: