Official

D - チーム選抜の重み合計 / Total Weight of Team Selection Editorial by admin

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\) です。

アルゴリズム

  1. \(W_i\) の総和 \(S = \sum_{i=1}^{N} W_i\) を計算する。
  2. 階乗とその逆元を前計算し、二項係数 \(\binom{n}{r}\) を \(O(1)\) で求められるようにする。
  3. \(\text{count} = \sum_{j=\lfloor N/2 \rfloor}^{N-1} \binom{N-1}{j}\) をループで計算する(ループ回数は約 \(N/2\))。
  4. 答えは \(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 によって生成されました。