D - チーム選抜の重み合計 / Total Weight of Team Selection 解説 by admin
Qwen3-Coder-480B概要
\(N\) 人の選手中から \(\lfloor N/2 \rfloor + 1\) 人以上の選手を選ぶすべての組み合わせについて、選ばれた選手の実力値の合計(スコア)を求め、その総和を \(998244353\) で割った余りを答える問題。
考察
この問題では、すべての「有効な選び方」(選ぶ人数が \(\lfloor N/2 \rfloor + 1\) 以上)におけるスコアの合計を求めたいです。素朴な方法としては、すべての部分集合を列挙して、条件を満たすものだけスコアを計算するという方法がありますが、これは \(2^N\) の組み合わせがあり、\(N\) が最大 \(5 \times 10^5\) なので到底間に合いません。
そこで、各選手の実力値 \(W_i\) が最終的な答えにどれだけ寄与するかを考えます。つまり、選手 \(i\) が含まれる「有効なチーム」の個数を数え、それに応じて \(W_i\) を答えに加算すればよいのです。
選手 \(i\) を含むチームの作り方は、「選手 \(i\) を固定して、残り \(N-1\) 人から \(\lfloor N/2 \rfloor\) 人以上を選ぶ」ということです。これは組み合わせの計算で求められます。
具体的には、選手 \(i\) を含むチームで、全体で \(k\) 人のチームを選ぶとき、残りの \(k-1\) 人は \(N-1\) 人から選ぶので、組み合わせの数は \(\binom{N-1}{k-1}\) です。そして、\(k\) は \(\lfloor N/2 \rfloor + 1\) から \(N\) まで動きます。
したがって、選手 \(i\) がスコアに寄与する回数は以下のように求められます:
\[ \sum_{k = \lfloor N/2 \rfloor + 1}^{N} \binom{N-1}{k-1} = \sum_{s = \lfloor N/2 \rfloor}^{N-1} \binom{N-1}{s} \]
この累積和を前計算しておけば、各選手について定数時間で寄与回数を求めることができます。
アルゴリズム
- まず、\(\binom{n}{k} \bmod 998244353\) を高速に求めるために、階乗とその逆元を前計算します。
- 各選手の寄与回数を求めるために、後ろから累積和を取った配列
suffを用意します: $\( \text{suff}[s] = \sum_{k=s}^{N-1} \binom{N-1}{k} \)$ - 各選手 \(i\) について、その実力値 \(W_i\) に寄与回数(
suff[min_select - 1])を掛けたものを答えに加算します。
計算量
- 時間計算量: \(O(N)\)
- 空間計算量: \(O(N)\)
実装のポイント
組み合わせの計算にはフェルマーの小定理による逆元計算が必要(\(998244353\) は素数)。
逆元の前計算は、後ろから DP 的に計算することで効率的に行えます。
答えは常に \(\bmod 998244353\) を取りながら計算しましょう。
ソースコード
def main():
MOD = 998244353
N = int(input())
W = list(map(int, input().split()))
# 前計算: 組み合わせ nCk mod MOD
# fact[i] = i! % MOD
# ifact[i] = (i!)^{-1} mod MOD
max_n = N
fact = [1] * (max_n + 1)
for i in range(1, max_n + 1):
fact[i] = fact[i-1] * i % MOD
def modinv(a):
return pow(a, MOD-2, MOD)
ifact = [1] * (max_n + 1)
ifact[max_n] = modinv(fact[max_n])
for i in range(max_n - 1, -1, -1):
ifact[i] = ifact[i+1] * (i+1) % MOD
def comb(n, k):
if k < 0 or k > n:
return 0
return fact[n] * ifact[k] % MOD * ifact[n-k] % MOD
min_select = N // 2 + 1 # この人数以上選ばなければならない
total = 0
# 各選手 i が何回スコアに寄与するかを数える
# それは、選手 i を含む有効なチームの数に等しい
# 選手 i を含むチームで、残り N-1 人から size-1 人を選ぶ
# size は min_select 〜 N
# つまり sum_{s=min_select-1}^{N-1} C(N-1, s)
# 事前に累積和を計算しておく
# suff[s] = sum_{k=s}^{N-1} C(N-1, k)
suff = [0] * (N + 1)
tmp = 0
for s in range(N-1, min_select - 2, -1):
tmp = (tmp + comb(N-1, s)) % MOD
suff[s] = tmp
for i in range(N):
w = W[i]
count = suff[min_select - 1]
total = (total + w * count) % MOD
print(total)
if __name__ == "__main__":
main()
この解説は qwen3-coder-480b によって生成されました。
投稿日時:
最終更新: