D - チーム選抜の重み合計 / Total Weight of Team Selection Editorial by admin
GPT 5.2 High概要
「人数が \(\lfloor N/2 \rfloor + 1\) 人以上」の部分集合(チーム)すべてについて、選ばれた実力値の合計(スコア)をさらに全て足し合わせた値を、\(998244353\) で割った余りで求めます。
考察
重要な気づき:各選手の寄与を数え上げる(線形性)
求めたい量は - 各有効な選び方 \(S\)(部分集合)についてのスコア \(\sum_{i \in S} W_i\) - それを全ての有効な \(S\) で合計
なので、 [ \sum{\text{valid } S}\ \sum{i\in S} Wi ] となります。ここで和の順序を入れ替えると [ \sum{i=1}^{N} W_i \cdot (\text{選手 }i\text{ を含む有効な選び方の数}) ] になります。つまり「選手 \(i\) が選ばれる有効な選び方が何通りあるか」だけ分かれば、全体が計算できます。
対称性:選手によらず個数は同じ
条件は「選ばれた人数」だけで決まるため、どの選手 \(i\) でも - 「\(i\) を含む有効な選び方の数」は同じ
よって、\(i\) ごとの個数をわざわざ別々に数える必要はなく、 [ (\sum_i W_i)\times \text{count} ] で済みます(count は「ある固定選手を必ず選ぶ」としたときの有効パターン数)。
素朴解が無理な理由
\(2^N\) 通りの全探索は、\(N \le 5\times 10^5\) では不可能です。
どう数えるか:固定選手を入れて残りから選ぶ
ある選手 \(i\) を必ず選ぶとします。残りは \(N-1\) 人(これを \(n=N-1\) とおく)で、そこから \(s\) 人選ぶと全体人数は \(1+s\) 人です。
有効条件: [ 1+s \ge \left\lfloor \frac{N}{2} \right\rfloor + 1 \quad\Longleftrightarrow\quad s \ge \left\lfloor \frac{N}{2} \right\rfloor ]
よって count は [ \sum_{s=\lfloor N/2\rfloor}^{n} \binom{n}{s} ] です。
例:\(N=5\) のとき有効は \(3\) 人以上。固定選手 1 人を含むので残り 4 人から \(2\) 人以上選ぶ: [ \binom{4}{2}+\binom{4}{3}+\binom{4}{4}=6+4+1=11 ] したがって答えは \((W_1+\cdots+W_5)\times 11\) になります。
アルゴリズム
- \(sumW = \sum_{i=1}^{N} W_i \bmod MOD\) を計算
- \(n=N-1\) とする(固定選手以外の人数)
- 二項係数を高速に出すため、階乗 \(fact[k]=k!\) と逆階乗 \(invfact[k]=(k!)^{-1}\) を \(k=0..N\) まで前計算する
- MOD は素数なので、逆元はフェルマーの小定理で \(a^{-1}\equiv a^{MOD-2}\)
- \(\binom{n}{r}=fact[n]\cdot invfact[r]\cdot invfact[n-r] \bmod MOD\)
- count を直接 \(\sum_{s=\lfloor N/2\rfloor}^{n}\binom{n}{s}\) と足してもよいですが、コードでは次を使います:
[
\sum{s=\lfloor N/2\rfloor}^{n}\binom{n}{s}
= 2^n - \sum{s=0}^{\lfloor N/2\rfloor-1}\binom{n}{s}
]
- \(2^n\) は「残り \(n\) 人の選び方総数」
- そこから「人数が足りない」ケースを引く
- 最終解: [ ans = sumW \times count \bmod MOD ]
計算量
- 時間計算量: \(O(N)\)
(階乗・逆階乗の前計算が \(O(N)\)、二項係数の和がだいたい \(O(N/2)\)) - 空間計算量: \(O(N)\)
(\(fact, invfact\) 配列)
実装のポイント
\(MOD=998244353\) は素数なので、逆元計算に
pow(x, MOD-2, MOD)が使えます。count は
2^n - prefixの形になるため、負にならないように%(MOD)を取って調整します。\(N=1\) のような小さいケースでも破綻しないように、\(\lfloor N/2\rfloor-1\) が負になる場合(コードでは
m=-1)は prefix を足さない処理にしています。ソースコード
import sys
MOD = 998244353
input = sys.stdin.readline
N = int(input())
W = list(map(int, input().split()))
sumW = sum(W) % MOD
n = N - 1 # choose from remaining players when one fixed player is included
# factorials up to N
fact = [1] * (N + 1)
for i in range(1, N + 1):
fact[i] = fact[i - 1] * i % MOD
invfact = [1] * (N + 1)
invfact[N] = pow(fact[N], MOD - 2, MOD)
for i in range(N, 0, -1):
invfact[i - 1] = invfact[i] * i % MOD
def comb(nn, rr):
if rr < 0 or rr > nn:
return 0
return fact[nn] * invfact[rr] % MOD * invfact[nn - rr] % MOD
# threshold: subset size >= floor(N/2)+1
# with fixed element included, choose s >= (floor(N/2)+1)-1 = floor(N/2) from n elements
# count = sum_{s=floor(N/2)}^{n} C(n,s) = 2^n - sum_{s=0}^{floor(N/2)-1} C(n,s)
m = N // 2 - 1 # upper index for excluded prefix: floor(N/2)-1
prefix = 0
if m >= 0:
if m > n:
m = n
for r in range(m + 1):
prefix += comb(n, r)
if prefix >= MOD:
prefix -= MOD
count = (pow(2, n, MOD) - prefix) % MOD
ans = sumW * count % MOD
print(ans)
この解説は gpt-5.2-high によって生成されました。
posted:
last update: