公式

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} \]

この累積和を前計算しておけば、各選手について定数時間で寄与回数を求めることができます。

アルゴリズム

  1. まず、\(\binom{n}{k} \bmod 998244353\) を高速に求めるために、階乗とその逆元を前計算します。
  2. 各選手の寄与回数を求めるために、後ろから累積和を取った配列 suff を用意します: $\( \text{suff}[s] = \sum_{k=s}^{N-1} \binom{N-1}{k} \)$
  3. 各選手 \(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 によって生成されました。

投稿日時:
最終更新: