公式

D - チーム選抜の重み合計 / Total Weight of Team Selection 解説 by kyopro_friends


有効な選び方であるための最少人数 \(\left\lfloor\frac{N}{2}\right\rfloor+1\)\(K\) とおきます。

各選手が答えに何度寄与するかを考えることで、求める答えは \(\sum_i W_i\times(選手 i を含む有効な選び方の数)\) となります。

選手 \(i\) を含む有効な選び方の数は、選手 \(i\) を選び、かつ、選手 \(i\) 以外から \(K-1\) 人以上選ぶ方法なので、\(\sum_{k=K-1}^{N-1}\binom{N-1}{k}\) となります。これは \(i\) によりません。よってこの値を \(C\) とおくと、求める答えは \(C\sum_i W_i\) となります。

二項係数は階乗及びその逆元を前計算しておくことで \(O(1)\) で求められるとしてよいため、 \(C\)\(O(N)\) で求めることができます。

よって前計算を含めこの問題を \(O(N+\log \mathrm{MOD})\) で解くことができました。

実装例 (C++)

#include<bits/stdc++.h>
#include<atcoder/modint>
using namespace std;
using modint=atcoder::modint998244353;

int main(){
  int n;
  cin >> n;
  vector<int>w(n);
  for(int i=0; i<n; i++) cin >> w[i];

  vector<modint>fact(n+1), invfact(n+1);
  fact[0]=1;
  for(int i=1; i<=n; i++){
    fact[i] = fact[i-1] * i;
  }
  invfact[n] = 1 / fact[n];
  for(int i=n-1; i>=0; i--){
    invfact[i] = invfact[i+1] * (i+1);
  }

  auto choose=[&](int n, int r){
    return fact[n] * invfact[r] * invfact[n-r];
  };

  modint c = 0;
  for(int k=n/2; k<n; k++){
    c += choose(n-1, k);
  }
  modint w_sum = 0;
  for(int i=0; i<n; i++){
    w_sum += w[i];
  }

  cout << (c * w_sum).val() << endl;
}

実装例 (Python)

MOD = 998244353

N = int(input())
W = list(map(int, input().split()))

fact = [0] * (N + 1)
fact[0] = 1
for i in range(1, N + 1):
    fact[i] = fact[i - 1] * i % MOD

invfact = [0] * (N + 1)
invfact[N] = pow(fact[N], -1, MOD)
for i in range(N - 1, -1, -1):
    invfact[i] = invfact[i + 1] * (i + 1) % MOD

def choose(n, r):
    return fact[n] * invfact[r] * invfact[n - r] % MOD

c = sum(choose(N - 1, k) for k in range(N // 2, N))

print(c * sum(W) % MOD)

なお、 \(\binom{n}{r}=\binom{n}{n-r}\)\(\sum_{k=0}^{n}\binom{n}{k}=2^n\) を用いることで、

\(C=\begin{cases} 2^{N-2} & N が偶数のとき\\ \frac{2^{N-1}+\binom{N-1}{K-1}}{2} & N が奇数のとき \end{cases}\)

となることが示せるため、答えを求めるために必要な二項係数の計算を高々 \(1\) 個にすることができ、階乗の前計算を不要にすることもできます。

MOD = 998244353
N = int(input())
W = list(map(int, input().split()))

def fact(n):
  crr = 1
  for i in range(1, n+1):
    crr = crr * i % MOD
  return crr

def choose(n, r):
  return fact(n) * pow(fact(r)*fact(n-r), -1, MOD) % MOD

if N % 2 == 0:
  c = pow(2, N-2, MOD)
else:
  c = (pow(2, N-1, MOD) + choose(N-1, N//2)) * ((MOD+1)//2) % MOD
print(sum(W) * c % MOD)

投稿日時:
最終更新: