CSAcademy Round #9 : E. Sum of Squares

TLが妙に厳しい。
https://csacademy.com/contest/round-9/#task/sum-of-squares

問題

整数N,Kが与えられる。
K要素の正整数からなる集合(multiset)のうち、総和がNになるものに対して各要素の２乗和を求めよ。

解法

K要素で総和がNとなる単調増加な数列Aを考える。
このような単調増加な数列の数え上げは、下記の問題が参考になる。
yukicoder : No.269 見栄っ張りの募金活動 - kmjp's blog

以下の状態を考える。
dp[i][j] := Aにおいて後ろからi個目までの要素の和がj(それ以前の要素は0)となる数列に対し{その組み合わせ数, ２乗和}

後ろからi個目より後ろ側の要素を1加算することを考える。(この処理はAの単調増加性を壊さない)
すなわちA' = 先頭(N-i)要素は0。A'[N-i]～A'[N-1]はそれぞれA[N-i]～A[N-1]に1加算したもの、とする。
sum(A'[i]^2 - A[i]^2) = 2 * sum(A) + iとなり、Aの和だけ考えればよく、Aの内訳を考える必要がない。

上記の前提は、Aの前の方が0であることである。
そこで、Aの各要素は最後に1インクリメントすることとし、AをK要素の"非負"整数からなる総和が(N-K)となる単調増加な数列となるようにしよう。
するとdp[i][j]は以下のように処理できる。
dp[i][j].first = dp[i-1][j].first + dp[i][j-i].first
dp[i][j].second = dp[i-1][j].second + (dp[i][j-i].second +dp[i][j-i].first*(2*(j-i)+i))

dp[K][N-K].secondは総和が(N-K)となる数列Aの２乗和なので、解は全体を1インクリメントしてdp[K][N].second + dp[K][N].first * (2*N+K)

int N,K;
ll mo=1000000007;

ll from[10101];
ll fromS[10101];

void add(ll& a,ll b) {
	a+=b;
	if(a>=mo) a-=mo;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	N-=K;
	from[0]=1;
	fromS[0]=0;
	
	for(i=1;i<=K;i++) {
		for(j=0;j+i<=N;j++) if(j+i<=N) {
			add(from[j+i],from[j]);
			add(fromS[j+i],(fromS[j]+from[j]*(2*j+i))%mo);
		}
	}
	cout<<(fromS[N]+from[N]*(2*N+K))%mo<<endl;
	
}

まとめ

自作問題が頭をよぎりすぎて、前から埋めることにこだわってしまったのが敗因。