公式

C - ボーナスの分配 / Bonus Distribution 解説 by admin

Claude 4.6 Opus (Thinking)

概要

\(N\) 人の社員が \(D\) ヶ月にわたって受け取るボーナス総額の積 \(R\) を、modular arithmetic(mod \(998244353\))を用いて求める問題です。式を整理すると、単純な累乗と積の計算に帰着できます。

考察

式の整理が鍵

各社員 \(i\) のボーナス総額は次の通りです:

\[T_i = \frac{P_i}{S} \times \sum_{j=1}^{D} W_j\]

ここで \(S = \sum_{i=1}^{N} P_i\) です。求めたいのは全社員のボーナス総額の積:

\[R = \prod_{i=1}^{N} T_i = \prod_{i=1}^{N} \left( \frac{P_i}{S} \times \sum_{j=1}^{D} W_j \right)\]

\(T_i\) の中で \(\frac{1}{S}\)\(\sum W_j\)\(i\) に依存しない共通因子なので、\(N\) 個の積をとると:

\[R = \frac{\left(\prod_{i=1}^{N} P_i\right) \times \left(\sum_{j=1}^{D} W_j\right)^N}{S^N}\]

このように、複雑に見える問題が 3つの値の計算 に帰着されます:

  1. \(\prod_{i=1}^{N} P_i\):全社員の貢献度の積
  2. \(\left(\sum_{j=1}^{D} W_j\right)^N\):ボーナス原資の合計の \(N\)
  3. \(S^N\):貢献度の合計の \(N\)

素朴に有理数のまま計算すると?

\(P_i\)\(W_j\) が最大 \(10^9\)\(N\) が最大 \(10^6\) なので、\(\prod P_i\) は桁数が膨大になり、多倍長整数で直接計算すると TLE になります。すべて \(\mod 998244353\) 上で計算するのがポイントです。

\(R = 0\) となるケース

\(P_i \geq 1\) なので \(\prod P_i > 0\) かつ \(S > 0\) です。\(R = 0\) となるのは \(\sum W_j = 0\) のとき、すなわち全ての \(W_j = 0\) のときだけです(\(W_j \geq 0\) なので)。このケースは先に判定して処理します。

アルゴリズム

  1. 入力を読み込む。
  2. \(\sum W_j = 0\)(全 \(W_j = 0\))なら答えは \(0\) を出力して終了。
  3. 以下をすべて \(\mod 998244353\) で計算する:
    • \(\text{prodP} = \prod_{i=1}^{N} P_i \mod M\)
    • \(\text{sumW} = \sum_{j=1}^{D} W_j \mod M\)
    • \(S = \sum_{i=1}^{N} P_i \mod M\)
  4. フェルマーの小定理を用いて:
    • \(\text{sumW}^N \mod M\)(繰り返し二乗法)
    • \(S^N \mod M\) の逆元 \((S^N)^{-1} \equiv (S^N)^{M-2} \mod M\)
  5. 答えは \(\text{prodP} \times \text{sumW}^N \times (S^N)^{-1} \mod M\)

具体例

\(N=2, D=2, P=[1,2], W=[3,6]\) のとき: - \(S = 3\), \(\sum W = 9\) - \(T_1 = \frac{1}{3} \times 9 = 3\), \(T_2 = \frac{2}{3} \times 9 = 6\) - \(R = 3 \times 6 = 18\) - 式で確認:\(R = \frac{(1 \times 2) \times 9^2}{3^2} = \frac{2 \times 81}{9} = 18\)

計算量

  • 時間計算量: \(O(N + D + \log M)\)(積・和の計算が \(O(N+D)\)、繰り返し二乗法が \(O(\log M)\)
  • 空間計算量: \(O(N + D)\)(入力の格納)

実装のポイント

  • \(\sum W_j = 0\) の判定は実際の値で行う:mod をとった後の値が \(0\) でも、実際の合計が \(0\) とは限りません。\(W_j \geq 0\) なので「全て \(0\) かどうか」で判定します。

  • \(S\)\(998244353\) の倍数でないことが保証されているため、\(S^N\) のモジュラ逆元が必ず存在します。

  • Python の pow(a, b, MOD) は繰り返し二乗法を内部で行うため、効率的にべき乗とモジュラ逆元を計算できます。

  • 大量入力に対応するため sys.stdin.buffer.read() で一括読み込みしています。

    ソースコード

import sys

def main():
    input_data = sys.stdin.buffer.read().split()
    idx = 0
    N = int(input_data[idx]); idx += 1
    D = int(input_data[idx]); idx += 1
    
    MOD = 998244353
    
    P = [int(input_data[idx + i]) for i in range(N)]
    idx += N
    W = [int(input_data[idx + i]) for i in range(D)]
    idx += D
    
    # T_i = P_i / S * sum(W)
    # R = prod(T_i) = prod(P_i) * (sum(W))^N / S^N
    
    # If any W_j causes sum(W) = 0, then R = 0 if sum(W)=0
    # Actually if sum(W) = 0 then R = 0
    
    sumW = sum(W) % MOD
    
    # If sumW == 0 mod something... but we need to check if actual sum is 0
    # sum(W) can be 0 only if all W_j = 0, since W_j >= 0
    # Actually sum of non-negative integers is 0 iff all are 0
    # So check if all W are 0
    
    actual_sumW_zero = all(w == 0 for w in W)
    
    if actual_sumW_zero:
        print(0)
        return
    
    # Also if any P_i = 0... but constraint says P_i >= 1, so prod(P_i) > 0
    
    # R = prod(P_i) * sumW^N / S^N
    # Compute everything mod MOD
    
    # S = sum(P_i)
    S = 0
    for p in P:
        S = (S + p) % MOD
    
    # prod(P_i) mod MOD
    prodP = 1
    for p in P:
        prodP = prodP * (p % MOD) % MOD
    
    # sumW^N mod MOD
    sumW_N = pow(sumW, N, MOD)
    
    # S^N mod MOD
    S_N = pow(S, N, MOD)
    
    # S_N inverse
    S_N_inv = pow(S_N, MOD - 2, MOD)
    
    ans = prodP * sumW_N % MOD * S_N_inv % MOD
    
    print(ans)

main()

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

投稿日時:
最終更新: