1 / 3
Oct 2023

I am solving TDKPRIME - Finding the Kth Prime
I am using Sieve of Eratosthenes to answer the K queries. The length of the sieve array is 86028122 which is just enough to store the 5000000th prime number which is 86028121.
I am getting TLE error. Below is my code. Please suggest to me how can I optimize it further.

class Main {
public static void main(String[] args) {
Scanner sc = new Scanner(System.in3);
int t = sc.nextInt();

    int SIEVE_LEN = 86028122;
    int PRIMES_LEN = 5000000;

    boolean[] sieve = new boolean[SIEVE_LEN];
    int[] primes = new int[PRIMES_LEN];
    int count = 0;

    sieve[0] = true; sieve[1] = true;

    // if prime -> sieve[i] = false, else -> sieve[i] = true
    for(int i = 2; i < SIEVE_LEN; i++) {
        if(sieve[i] == false) {
            primes[count] = i;
            count++;
            for(int j = 2*i; j < SIEVE_LEN; j+=i) {
                sieve[j] = true;
            }
        }
    }

    while(t-- > 0) {
        int n = sc.nextInt();
        System.out.println(primes[n-1]);
    }
}

}

  • created

    Oct '23
  • last reply

    Dec '23
  • 2

    replies

  • 677

    views

  • 3

    users

  • 2

    links

Testing on IDEONE shows that the code takes around 1.7s to run. You can make it a little bit quicker by changing

for(int j = 2*i; j < SIEVE_LEN; j+=i)

to

for(int j = i*i; j < SIEVE_LEN; j+=i)

but then you’ll also need to prevent overflows of the i*i. On IDEONE, this got the time down to 1.2s, so that might just be enough. Failing that, you’ll need to use a smarter sieve. Or a different language.

2 months later

You do not have to track sieve for even numbers. Keeping info only for odd numbers will decrease the number of operations twice, as well as the memory size, that may be even more important.
You can use int bit array instead of boolean for sieve, so each element will require 1 bit instead of 1 byte. It looks like a bit more of processor work, but decrease memory usage 8 more times, that is more important since the memory latency seems to be the bottleneck in this problem.
And, yes, you can start filling the sieve from j*j instead of 2*j, and when j*j becomes >= SIEVE_LEN, you can stop sieving and simply collect all the remaining primes.

Update: for me, the most important problem here was fast input and output.
My sieve version works in 0.36 seconds at ideone, input is hand-made and output if kind of “buffered”, so the problem is accepted in 0.48 seconds.
The main loop looks like:

InputStream in = System.in;
int t = readUint(in);
StringBuilder sb = new StringBuilder();
for (int it = 0; it < t; it++) {
    int n = readUint(in);
    sb.append(primes[n - 1]); 
    sb.append('\n');
}
System.out.print(sb.toString());

and readUint function looks like:

public static int readUint(InputStream in) throws java.io.IOException
{
	int u = 0;
	for(;;) {
		int c = in.read();
		assert (0 <= c);
		u = c - 48;
		if (0 <= u && u < 10)
			break;
	}
	for (;;) {
		int v = in.read() - 48;
		if (0 > v || v >= 10)
			break;
		u = u * 10 + v;
	}
	return u;
}

Probably, a hand-made version of output (in bytes, without chars) could save some milliseconds.
Creating a BufferedInputStream around System.in increased runtime to 0.52s, so System.in looks already buffered.