AOJ 1241 Lagrange's Four-Square Theorem

問題

Lagrange's Four-Square Theorem

四平方定理っていうのがある。全ての自然数は高々4個の平方数の和で表すことができる。nが与えられる。高々4個の平方数の和がnになる組み合わせの数を答える。

解法

全探索とdpで解けます。両方で解こうと思います

心臓に悪いコード

int n;

int dfs(int sum, int d, int c){
  if(sum == n) return 1;
  if(c >= 4) return 0;
  int res = 0;
  for(int i = d; i*i <= n; i++){
    if(sum + i*i > n) break;
    res += dfs(sum + i*i, i, c+1);
  }
  return res;
}

int main(){
  while(scanf("%d", &n) && n)
    printf("%d\n", dfs(0, 1, 0));
  return 0;
}

i*i <= n

の部分を初め

i <= sqrt(n)

にしてたらTLE食らった。

動的計画法で解く。

dp[i][j][k] := i番目までの平方数でj個使ってkが何個できるか