Codeforces #572 Div1 C. Array Beauty

計算量の見積もりが難しい。
http://codeforces.com/contest/1188/problem/C

問題

N個の整数A[i]が与えられる。
ここからK個を選んだ時、2要素の最小値をその選んだ集合のスコアとする。

K個の選び方全通りにおける、スコアの総和を求めよ。

解法

f(n) := スコアがn以上な組み合わせ数
g(n) := スコアがnちょうどな組み合わせ数
とすると、
g(n) = f(n)-f(n+1)となり、g(n)*nの総和が求める値となる。

f(n)は単純なDPで求められる。
一見計算量はO(max(A)*K)かかりそうだが、Kが大きい場合はnはmax(A)/(K-1)程度までしか取りえないのでTLEしない。

int N,K;
int A[1010];
ll pat[101010];

ll from[1010],S[1010],to[1010];
ll mo=998244353;

ll hoge(int d) {
	ZERO(from);
	ZERO(to);
	int i,j,x,y;
	FOR(i,N) from[i]=1;
	FOR(i,K-1) {
		S[0]=from[0];
		for(j=1;j<N;j++) (S[j]=S[j-1]+from[j])%=mo;
		int x=-1;
		FOR(y,N) {
			while(x<y && A[y]-A[x+1]>=d) x++;
			if(x==-1) to[y]=0;
			else to[y]=S[x];
		}
		swap(from,to);
	}
	ll sum=0;
	FOR(i,N) sum+=from[i];
	return sum%mo;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	FOR(i,N) cin>>A[i];
	sort(A,A+N);
	for(i=1;i<=100001/(K-1);i++) {
		pat[i]=hoge(i);
	}
	ll ret=0;
	for(i=1;i<=100001/(K-1);i++) {
		ret+=i*(pat[i]-pat[i+1]+mo)%mo;
	}
	cout<<ret%mo<<endl;
	
}

まとめ

計算量の見積もりをまんまとミスした。