AT_abc471_e

AtCoder - abc471_e

闲得蛋疼写片

题解

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int maxn=2e6+5;
const ll mod=998244353;


ll fac[maxn],ifac[maxn];

ll powmod(ll x,ll y) {
	if (!y) return 1;
	ll t=powmod(x,y/2);
	if (y&1) return t*t%mod*x%mod;
	else return t*t%mod;
}

ll C(ll n,ll m) {
	
	if (m<0||n<0||n<m) return 0;
	return fac[n]*ifac[m]%mod*ifac[n-m]%mod;
}
ll n,k;
void solve() {
    cin>>n>>k;
    ll s1=0,s2=0;
    for (int i=1;i<=n;i++) {
    	ll x;cin>>x;
    	s1=(s1+x)%mod;
    	s2=(s2+x*x%mod)%mod;
	}
	ll ans=((s1*s1%mod*C(n-2,k-2)%mod-s2*C(n-2,k-2)%mod+mod)%mod+s2*C(n-1,k-1)%mod)%mod;
	cout<<ans<<'\n';
}

void prework() {
	fac[0]=1;
	for (ll i=1;i<maxn;i++) fac[i]=fac[i-1]*i%mod;
	ifac[maxn-1]=powmod(fac[maxn-1],mod-2);
	for (ll i=maxn-2;i>=0;i--) ifac[i]=ifac[i+1]*(i+1)%mod;
}

int main() {
    ios::sync_with_stdio(0);
    prework();
    int t=1;
    while (t--) solve();
}

\[\displaystyle(\sum_{x\in{S}} x)^{2}\\ =(x_{1}+x_{2}+...+x_{k})·(x_{1}+x_{2}+...+x_{k})+...+(x_{n-k+1}+x_{n-k+2}+...+x_{n})·(x_{n-k+1}+x_{n-k+2}+...+x_{n}) \]

显然最终可以变成一个式子：（其中A，B是系数）

\[A \times \displaystyle\sum_{i=1}^{n} {x_{i}^{2}} + B \times \displaystyle\sum_{i,j\in[1,n],i\neq j}{x_ix_j} \]

考虑每个数自己和自己相乘的次数

即在确定了一个数的情况下，在剩下的 \(n-1\) 个数中再选 \(k-1\) 个数的方案：\(A=C^{k-1}_{n-1}\)

考虑两个每个数互相乘的次数

即在确定了两个数的情况下，在剩下的 \(n-2\) 个数中再选 \(k-2\) 个数的方案：\(B=C^{k-2}_{n-2}\)

所以答案就是

\[C^{k-1}_{n-1} \times \displaystyle\sum_{i=1}^{n} {x_{i}^{2}} + C^{k-2}_{n-2} \times \displaystyle\sum_{i,j\in[1,n],i\neq j}{x_ix_j}\\ =C^{k-1}_{n-1} \times \displaystyle\sum_{i=1}^{n} {x_{i}^{2}}+C^{n-2}_{k-2}\times (\displaystyle\sum^{n}_{i=1}{x})^{2}-C^{n-2}_{k-2} \times \displaystyle\sum_{i=1}^{n} {x_{i}^{2}} \]

根据初中数学知识转化一下式子就可以O(n) 做了