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つの値の計算 に帰着されます:
- \(\prod_{i=1}^{N} P_i\):全社員の貢献度の積
- \(\left(\sum_{j=1}^{D} W_j\right)^N\):ボーナス原資の合計の \(N\) 乗
- \(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\) なので)。このケースは先に判定して処理します。
アルゴリズム
- 入力を読み込む。
- \(\sum W_j = 0\)(全 \(W_j = 0\))なら答えは \(0\) を出力して終了。
- 以下をすべて \(\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\)
- フェルマーの小定理を用いて:
- \(\text{sumW}^N \mod M\)(繰り返し二乗法)
- \(S^N \mod M\) の逆元 \((S^N)^{-1} \equiv (S^N)^{M-2} \mod M\)
- 答えは \(\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 によって生成されました。
投稿日時:
最終更新: