Perfect squares

Given a positive integer n, find the least number of perfect square numbers (for example, 1, 4, 9, 16, ...) which sum to n.

For example, given n = 12, return 3 because 12 = 4 + 4 + 4; given n = 13, return 2 because 13 = 4 + 9.

This question utilizes Lagrange's four-square theorem, which indicates all numbers can be represented by the sum of at most four square numbers. So we know the answer is 1, 2, 3 and 4. If the number mod 8 equals 7, it will be represented by 4 numbers (don't ask me why, ask Google). Starting from 0, check if a * a + b * b == n, it true. Check if one of a or b is 0, if true, return 1, else return 2. If we cannot find such a number, return 3.

We can further optimize the problem by divide all factors of 4 and 9.

public int numSquares(int n) {
        while (n % 4 == 0) n /= 4;
        while (n % 9 == 0) n /= 9;
        if (n % 8 == 7)
            return 4;
        for (int a = 0; a * a <= n; a++) {
            int b = (int)Math.sqrt(n - a * a);
            if (a * a + b * b == n) {
                if (a == 0 || b == 0)
                    return 1;
                else
                    return 2;
            }
        }
        return 3;
    }