Primes & Squares

Question

Codeforces #230B

Question

Whether a number has exactly three factors. (i.e. 1, itself and its square root, thus squares of primes) [Time limit is 2 sec, Memory 256 MB]

Input

Number of testcases, n (~10⁵) then numbers (~10¹²)

Output

Yes or No

Problem

I first check if it is a perfect square, then check if it's root is prime or not. Earlier I was trying with naive prime checking algorithm, but now I thought to speed up I should use sieve of Eratosthenes, which wouldn't take much time as largest number is 10¹² so square root is at max 10⁶ which is not a very large number for \$O(n \log \log n)\$ (778 microsec considering 1GHz ~\$10^9\$ ops/sec and 6000 microsec for \$O(n \log n)\$).

My sieve is a little bit different: (i) I have only indexed odd number to half the memory required. (0→3, 1→5, 2→7, …, x→2x+3) (ii) I use inNotPrime instead of isPrime because default value of boolean array is false, so it would be easy to avoid Arrays.fill(isPrime,true).

import java.util.Scanner;

public class Solution {
    public static void main(String[] args) {
    Scanner sc = new Scanner(System.in);
    int n = sc.nextInt();
    seivePrimes();
    for (int i = 0; i < n; i++) {
        long a = sc.nextLong();
        long sqrt = (int) Math.sqrt(a);
        if (a >= 4 && sqrt * sqrt == a) {
        if (isPrime((int) sqrt)) {
            System.out.println("YES");
        } else {
            System.out.println("NO");
        }
        } else {
        System.out.println("NO");
        }
    }
    sc.close();
    }

    static int maxlim = 1000000;
    private static void seivePrimes() {
    for (int i = 3; i * i <= maxlim; i += 2) {
        if (!isNotPrime[(i - 3) / 2]) {
        for (int j = i; i * j <= maxlim; j += 2) {
            isNotPrime[(i * j - 3) / 2] = true;
        }
        }
    }
    }

    static boolean[] isNotPrime = new boolean[(maxlim - 3) / 2 + 1];//isNotPrimes[x]=>2*x+3 is not prime
    private static boolean isPrime(int t) {
    if (t == 2)
        return true;
    if (t % 2 == 0)
        return false;
    else
        return !isNotPrime[(t - 3) / 2];
    }
}

What is making it slow so as that time limit is exceeded? I saw the solution to the problem afterwards and even they are using the same thing, squares and sieving.

Answer 1

What is making it slow so as that time limit is exceeded

Funnily enough, it can be the println. It's synchronized and (what's worse IIRC) it flushes the output. Accumulating it in a StringBuilder can be the solution (an answer of mine claiming this on SO was accepted).

Note that division and modulus are slow. a % 2 and a / 2 can't be in general optimized to a & 1 and a >> 2, respectively. The JIT can't do this when it can't see that a >= 0 holds. You can.

    if (a >= 4 && sqrt * sqrt == a) {
    if (isPrime((int) sqrt)) {
        System.out.println("YES");
    } else {
        System.out.println("NO");
    }
    } else {
    System.out.println("NO");
    }

That's bad! Always separate computation from IO. Write a method. When doing this, combine the conditions:

boolean isSquaredPrime(long a) {
    long sqrt = (int) Math.sqrt(a);
    return a >= 4 && sqrt * sqrt == a && isPrime((int) sqrt);
}

Note that it fails for a above \$2^{62}\$.

---

This part is fine (though I'd recommend against leaving out braces)

private static boolean isPrime(int t) {
if (t == 2)
    return true;
if (t % 2 == 0)
    return false;
else
    return !isNotPrime[(t - 3) / 2];
}

but it can be a bit compact, e.g.,

private static boolean isPrime(int t) {
    if ((t & 1) == 0) { // optimized evenness test
        return t == 2; // the only even prime
    } else {
        return !isNotPrime[(t - 3) >> 1];
    }
}

I wrote "compact" rather than "readable" as it's bad for people unaware of bit operations. But once you get used to them....

---

Method extraction is good for readability and reusability. It may sometimes make the program faster as shorter methods can be better optimized. The method call cost may be substantial, but then the method gets most probably inlined (which sort of undoes the method extraction, but readability gains remain).