D - チーム選抜の重み合計 / Total Weight of Team Selection Editorial by admin Editorial for printing
Claude 4.6 Opus (Thinking)概要
\(N\) 人の選手から \(\lfloor N/2 \rfloor + 1\) 人以上を選ぶ全ての選び方について、選ばれた選手の実力値の合計(スコア)を求め、それらを全て足し合わせた値を \(998244353\) で割った余りで出力する問題です。
考察
素朴なアプローチとその問題点
全ての有効な選び方を列挙してスコアを計算する方法を考えると、部分集合の数は最大 \(2^N\) 通り(\(N \leq 5 \times 10^5\))なので、到底間に合いません。
重要な気づき:和の順序を入れ替える(寄与の分解)
求めたい値は次のように書けます:
\[\text{答え} = \sum_{\text{有効な選び方 } S} \sum_{i \in S} W_i\]
ここで和の順序を入れ替えます。つまり「各選び方ごとにスコアを計算して足す」のではなく、「各選手 \(i\) が全体にどれだけ寄与するか」を考えます:
\[\text{答え} = \sum_{i=1}^{N} W_i \times (\text{選手 } i \text{ を含む有効な選び方の数})\]
各選手の寄与の計算
選手 \(i\) を必ず選ぶと固定すると、残り \(N-1\) 人から追加で何人か選んで、合計人数が \(\lfloor N/2 \rfloor + 1\) 人以上になればよいです。合計 \(k\) 人選ぶとき、選手 \(i\) 以外に \(k-1\) 人を \(N-1\) 人から選ぶので:
\[\text{(選手 } i \text{ を含む有効な選び方の数)} = \sum_{k=\lfloor N/2 \rfloor + 1}^{N} \binom{N-1}{k-1} = \sum_{j=\lfloor N/2 \rfloor}^{N-1} \binom{N-1}{j}\]
この値は全ての選手で共通です(\(W_i\) の値に依存しない)。したがって:
\[\text{答え} = \left( \sum_{i=1}^{N} W_i \right) \times \sum_{j=\lfloor N/2 \rfloor}^{N-1} \binom{N-1}{j}\]
具体例(\(N=5\))
閾値は \(\lfloor 5/2 \rfloor + 1 = 3\) 人以上。各選手を含む有効な選び方の数は:
\[\sum_{j=2}^{4} \binom{4}{j} = \binom{4}{2} + \binom{4}{3} + \binom{4}{4} = 6 + 4 + 1 = 11\]
答えは \((W_1 + W_2 + W_3 + W_4 + W_5) \times 11 \mod 998244353\) です。
アルゴリズム
- \(W_i\) の総和 \(S = \sum_{i=1}^{N} W_i\) を計算する。
- 階乗とその逆元を前計算し、二項係数 \(\binom{n}{r}\) を \(O(1)\) で求められるようにする。
- \(\text{count} = \sum_{j=\lfloor N/2 \rfloor}^{N-1} \binom{N-1}{j}\) をループで計算する(ループ回数は約 \(N/2\))。
- 答えは \(S \times \text{count} \mod 998244353\)。
計算量
- 時間計算量: \(O(N)\)(階乗の前計算に \(O(N)\)、二項係数の合計計算に \(O(N/2)\))
- 空間計算量: \(O(N)\)(階乗テーブルと逆元テーブル)
実装のポイント
逆元の計算:\(998244353\) は素数なので、フェルマーの小定理により \(a^{-1} \equiv a^{p-2} \pmod{p}\) で逆元を求められます。階乗の逆元テーブルは \(N!\) の逆元から逆順に計算すると効率的です。
ループ範囲:二項係数を足すループは \(j = \lfloor N/2 \rfloor\) から \(N-1\) までで、約 \(N/2\) 回の反復なので \(N = 5 \times 10^5\) でも十分高速です。
全選手で寄与が同じという点に気づけるかが本問の核心です。寄与分解(linearity of expectation / summation の考え方)は競技プログラミングで頻出のテクニックです。
ソースコード
import sys
def solve():
MOD = 998244353
input_data = sys.stdin.buffer.read().split()
N = int(input_data[0])
W = [int(input_data[i+1]) for i in range(N)]
# We need to sum over all "valid" subsets S (|S| >= floor(N/2)+1) the sum of W_i for i in S.
#
# By linearity of summation, the answer = sum over i of W_i * (number of valid subsets containing i)
#
# For a fixed player i, the number of valid subsets containing i =
# sum_{k = floor(N/2)+1}^{N} C(N-1, k-1)
# because we fix player i in the subset, then choose k-1 more from the remaining N-1 players.
#
# Let threshold = floor(N/2) + 1
# count = sum_{k=threshold}^{N} C(N-1, k-1) = sum_{j=threshold-1}^{N-1} C(N-1, j)
#
# This count is the same for every player i.
# So answer = count * sum(W) mod MOD.
S = sum(W) % MOD
threshold = N // 2 + 1 # minimum team size
# count = sum_{j = threshold-1}^{N-1} C(N-1, j)
# Precompute factorials and inverse factorials
max_n = N
fact = [1] * (max_n + 1)
for i in range(1, max_n + 1):
fact[i] = fact[i-1] * i % MOD
inv_fact = [1] * (max_n + 1)
inv_fact[max_n] = pow(fact[max_n], MOD - 2, MOD)
for i in range(max_n - 1, -1, -1):
inv_fact[i] = inv_fact[i+1] * (i+1) % MOD
def comb(n, r):
if r < 0 or r > n:
return 0
return fact[n] * inv_fact[r] % MOD * inv_fact[n-r] % MOD
# count = sum_{j = threshold-1}^{N-1} C(N-1, j)
n1 = N - 1
lo = threshold - 1 # = floor(N/2)
count = 0
for j in range(lo, n1 + 1):
count = (count + comb(n1, j)) % MOD
# But this loop could be up to N/2 iterations which is up to 2.5*10^5, should be fine.
ans = S * count % MOD
print(ans)
solve()
この解説は claude4.6opus-thinking によって生成されました。
posted:
last update: