Official

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\) になります。

アルゴリズム

  1. \(sumW = \sum_{i=1}^{N} W_i \bmod MOD\) を計算
  2. \(n=N-1\) とする(固定選手以外の人数)
  3. 二項係数を高速に出すため、階乗 \(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\)
  4. 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\) 人の選び方総数」
    • そこから「人数が足りない」ケースを引く
  5. 最終解: [ 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: