Writing a number as sum of squares

[Saitama]

Here's the problem statement from HackerRank: https://www.hackerrank.com/contests/programaniacs-june-15/challenges/sum-of-squares-1

Problem Statement

Find the number of n-tuples $(a_1,a_2,a_3,\cdots, a_n)$ such that

$$a_1^2+a^2_{2}+a^2_{3}+\cdots+a^2_{n} = k\,\,\,\text{and}\,\,\, a_1≤a_2≤\cdots≤a_n\,\, \text{where} \,\,a_i>0$$

Input Format

The first line contains the number of test cases $T$.

The next $T$ lines contain two integers, $n$ and $k$.

Constraints

$1≤T≤10^5$
$1≤n≤100$
$1≤k≤5000$

Output Format

For each of the $T$ test cases, output one integer, the answer to the problem.

Sample Input

4
1 2
1 4
2 50
3 1000

Sample Output

0
1
2
2

Since the constraints are small, I tried a DP solution. Code I have written so far:

#include <cmath>
#include <cstdio>
#include <vector>
#include <iostream>
#include <algorithm>
using namespace std;

typedef long long ll;

ll dp[5010][101]={};

int main() {
    ll T, i, j, n, k;
    vector<ll> sq;
    for(i=1; i*i<=5000; ++i)
        sq.push_back(i*i);
    for(i=0; i<sq.size(); ++i)
        dp[sq[i]][1]=1;
    for(i=2; i<=5000; ++i)
    {
        for(j=2; j<=min(i, 100LL); ++j)
        {
            for(k=0; k<sq.size(); ++k)
            {
                if(sq[k]>i)
                    break;
                dp[i][j]=dp[i][j]+dp[i-sq[k]][j-1];
            }
        }
    }
    cin>>T;
    while(T--)
    {
        cin>>n>>k;
        cout<<dp[k][n]<<endl;
    }
    return 0;
}

dp[j] denotes the number of ways to write $i$ as sum of $j$ squares.

The problem with the above code that it overcounts the number of ways. How do I fix this issue?

Thanks!

 

[I like Serena]

Hi Pranav,

The typical reason to overcount in a situation like this, is that a solution with the same value more than once, will be counted more than once.

So for instance $1^2+1^2+2^2+2^2+3^2+3^2+3^2 = k$ will typically be counted $2! \cdot 2! \cdot 3!$ times.

Easiest way to fix it, is to track the solutions and eliminate duplicates.
Alternatively, we can compensate and count for instance my example only for the fraction $1 / (2! \cdot 2! \cdot 3!)$.

Unfortunately, I guess that doesn't work very well with your dynamic programming approach.
Alternatively, we could try a back tracking approach that will yield full solutions.

 

[Saitama]

Alternatively, we could try a back tracking approach that will yield full solutions.

But that won't run under the given time constraints. (Sweating)

I can't think of a way besides dp. :(