Isa@Diary

ソフトウェア開発やってます。プログラミングとか、US生活とかについて書きます。

AOJ 1241:Lagrange's Four-Square Theorem

AOJ上でソースを公開できるようになったのでそのテストも兼ねて

ソース
http://judge.u-aizu.ac.jp/onlinejudge/review.jsp?rid=376363

与えられたN(1<=N<=2^15)を4個以下の平方数の和で表す方法は何通りあるか、という問題。
2^15以下の平方数は181個しかないのでDPで

dp[i][j][k] := i番目の平方数まででj個使った和がkになるようなものの個数

として、i番目の平方数をα個使ったときに

dp[i][j+α][k+i*i*α] += dp[i-1][j][k];

としていけばよい。

答えは与えられたnについて
dp[181][0][n]からdp[181][4][n]までの和を取ってやればいい。
計算量も181*5*5*2^15で十分間に合う。

DP書けると楽しいけどいつももうちょっと早い漸化式の立て方があるのではと思ってしまう。