CSAcademy Round #9: Sum of Squares

https://csacademy.com/contest/round-9/#task/sum-of-squares

解法

$a_1+\cdots a_K=N$ となるような集合は N の K 分割である。分割数は次のように構成できるのであった。

• P[i][j] := j を i 個の整数に分割するパターン

(1) $0 \in \{a_0, \ldots, a_{i-1}\}$ のとき
盤面の高さを 1 減らしたものを見ればいい。P[i][j] += P[i-1][j]

(2) $0 \not\in \{a_0, \ldots, a_{i-1}\}$ のとき
一番左のタイルをすべて削ったものと一対一に対応づらけられる。P[i][j] += P[i][j-i]

この構成法を使えば次のような DP ができる。

• dp[i][j] := （j の i 分割の 2 乗和）の総和

今回は正整数なので、あらかじめ 1 を引いて非負整数の議論に持ち込もう。

まず (1) のパターンの遷移を考える。タイルの個数が変わってないので変化なしに見えるが、あらかじめ 1 を引いていることを考えると総和は 1 増えている。 よって全体では P[i-1][j] 増加する。

• $dp[i][j] \gets dp[i-1][j] + P[i-1][j]$

次に (2) のパターンの遷移を考える。j-iのi分割が $a_0,\ldots,a_{i-1}$ だったとすると、これに対応する j の i 分割は $a_0+1,\ldots,a_{i-1}+1$ である。 この二乗和の増分は次のように表せる。

$$ \sum_{k=0}^{i-1} \left(\left(a_k+1)^2-a_k^2 \right)\right)=2j+i $$

したがって

• $dp[i][j] \gets dp[i][j-i] + (2j+i)P[i][j-i]$

あとはこの漸化式にしたがって計算していけばいい。N,K≦10000 なので重そうだが、割と間に合う。

TL が 500 ms なの攻めすぎなのでは。