TLが妙に厳しい。
https://csacademy.com/contest/round-9/#task/sum-of-squares
問題
整数N,Kが与えられる。
K要素の正整数からなる集合(multiset)のうち、総和がNになるものに対して各要素の2乗和を求めよ。
解法
K要素で総和がNとなる単調増加な数列Aを考える。
このような単調増加な数列の数え上げは、下記の問題が参考になる。
yukicoder : No.269 見栄っ張りの募金活動 - kmjp's blog
以下の状態を考える。
dp[i][j] := Aにおいて後ろからi個目までの要素の和がj(それ以前の要素は0)となる数列に対し{その組み合わせ数, 2乗和}
後ろから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の2乗和なので、解は全体を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; }
まとめ
自作問題が頭をよぎりすぎて、前から埋めることにこだわってしまったのが敗因。