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)
投稿日時:
最終更新:
