class Solution {

    // O(sqrt(n)) (constrained by prime factorization)
    //
    // We use the four-square theorem to notice that the result is always at most 4.
    // We can then check whether n is a square to find out whether the result is 1, if not we can use two-square and
    // three-square theorems to find out whether the result is 2 or 3, else it's 4.
    //
    // If someone finds a way to check whether n is a sum of two squares efficiently, we can reduce the complexity
    // to O(log n)! (Same if the prime factorization is known)


    public static int numSquares(int n) {
        if (n == 0) return 0;
        if (n == (int) Math.pow(Math.round(Math.sqrt(n)), 2)) return 1;

        Map<Integer, Integer> primeFactors = findPrimeFactors(n);
        if (primeFactors.entrySet().stream().filter(e -> e.getKey() % 4 == 3).allMatch(e -> e.getValue() % 2 == 0)) return 2;
        if (primeFactors.getOrDefault(2, 0) % 2 == 0 && (n >>> Integer.numberOfTrailingZeros(n)) % 8 == 7) return 4;
        return 3;
    }





    public static Map<Integer, Integer> findPrimeFactors(int n) {
        Map<Integer, Integer> res = new HashMap<>();
        for (int i = 2; i*i <= n; i++) {
            while (n % i == 0) {
                res.merge(i, 1, Integer::sum);
                n /= i;
            }
        }
        if (n > 1) res.merge(n, 1, Integer::sum);
        return res;
    }
}