CF 230B - T-primes

We are asked to identify T-primes in a list of positive integers. A T-prime is defined as a number that has exactly three distinct positive divisors. Thinking about divisors, the only way a number can have exactly three is if it is the square of a prime number.

codeforcescompetitive-programmingbinary-searchimplementationmathnumber-theory

CF 230B - T-primes

Rating: 1300
Tags: binary search, implementation, math, number theory
Solve time: 55s
Verified: yes

Solution

Problem Understanding

We are asked to identify T-primes in a list of positive integers. A T-prime is defined as a number that has exactly three distinct positive divisors. Thinking about divisors, the only way a number can have exactly three is if it is the square of a prime number. For example, 4 is a T-prime because its divisors are 1, 2, and 4, and 2 is prime. A number like 6 has four divisors (1, 2, 3, 6), so it is not T-prime.

The input provides the number of integers, up to 100,000, followed by the integers themselves, which can be as large as 10^12. Because the input size is significant, any solution that inspects divisors one by one for each number will be far too slow. For example, checking all numbers up to √10^12 for divisibility takes roughly a million steps per number, which is feasible for a single number but not for 100,000 of them. Therefore, we need a smarter approach that reduces the work per query.

We also need to handle edge cases carefully. The number 1 is never T-prime because it has only one divisor. Large squares like 10^12 must be handled without integer overflow or precision loss when taking square roots. Numbers that are perfect squares but whose square roots are not prime, such as 16 (4^2), must also be rejected.

Approaches

The brute-force approach is straightforward: for each number, count its divisors. If the count is exactly three, print YES; otherwise, print NO. Counting divisors naïvely requires iterating up to √x, giving a worst-case complexity of O(n√x), which is roughly 10^11 operations in the worst case. This is too slow for the problem constraints.

The key observation is that a number has exactly three divisors if and only if it is a perfect square of a prime number. This reduces the problem to two steps: determine if a number is a perfect square and if its square root is prime. To handle large input efficiently, we can generate all primes up to 10^6 using the Sieve of Eratosthenes (since √10^12 = 10^6) and then store their squares in a set. For each input number, we simply check if it is in the set of T-primes. This reduces the per-query work to a simple set lookup, which is O(1) on average.

Approach Time Complexity Space Complexity Verdict
Brute Force O(n√x) O(1) Too slow
Optimal O(n + √max(x) log log √max(x)) O(√max(x)) Accepted

Algorithm Walkthrough

  1. Compute all prime numbers up to 10^6 using the Sieve of Eratosthenes. This range is sufficient because any T-prime cannot have a square root larger than 10^6.
  2. Initialize an empty set to store T-primes. For each prime number p generated by the sieve, compute p squared and add it to the set. This precomputes all possible T-primes up to 10^12.
  3. Read the array of numbers. For each number x, check if x is in the set of precomputed T-primes.
  4. If x is found in the set, print YES; otherwise, print NO.

The reason this works is that the set contains exactly all numbers that are squares of primes. Every T-prime is included because its square root is prime, and any number not in the set either is not a perfect square or has a non-prime square root.

Python Solution

import sys
input = sys.stdin.readline

def sieve(limit):
    is_prime = [True] * (limit + 1)
    is_prime[0] = is_prime[1] = False
    for i in range(2, int(limit ** 0.5) + 1):
        if is_prime[i]:
            for j in range(i*i, limit + 1, i):
                is_prime[j] = False
    primes = [i for i, val in enumerate(is_prime) if val]
    return primes

def main():
    n = int(input())
    numbers = list(map(int, input().split()))
    
    max_root = int(10**6) + 1
    primes = sieve(max_root)
    
    t_primes = set(p * p for p in primes)
    
    for x in numbers:
        if x in t_primes:
            print("YES")
        else:
            print("NO")

if __name__ == "__main__":
    main()

The sieve function generates primes efficiently and avoids checking divisibility repeatedly. Storing T-primes in a set allows constant-time membership checks, which is critical for 100,000 queries. The code reads input using sys.stdin.readline to avoid slow input handling in Python. Using integer multiplication avoids any floating-point inaccuracies.

Worked Examples

Sample 1 input:

3
4 5 6
x √x Is prime √x T-prime?
4 2 Yes YES
5 2.236 No NO
6 2.449 No NO

This demonstrates that only perfect squares of primes are accepted.

Another example:

5
1 9 16 25 49
x √x Is prime √x T-prime?
1 1 No NO
9 3 Yes YES
16 4 No NO
25 5 Yes YES
49 7 Yes YES

This confirms the algorithm correctly rejects perfect squares of non-primes and accepts perfect squares of primes.

Complexity Analysis

Measure Complexity Explanation
Time O(n + √max(x) log log √max(x)) Sieve takes √max(x) log log √max(x), set lookup for n numbers is O(n)
Space O(√max(x)) Store primes and their squares

With n up to 10^5 and max(x) up to 10^12, the algorithm easily fits in a 2-second limit.

Test Cases

import sys, io

def run(inp: str) -> str:
    sys.stdin = io.StringIO(inp)
    output = io.StringIO()
    sys.stdout = output
    main()
    return output.getvalue().strip()

# provided sample
assert run("3\n4 5 6\n") == "YES\nNO\nNO", "sample 1"

# minimum-size input
assert run("1\n1\n") == "NO", "minimum input"

# perfect squares of primes
assert run("4\n4 9 25 49\n") == "YES\nYES\nYES\nYES", "primes squares"

# perfect squares of non-primes
assert run("3\n16 36 100\n") == "NO\nNO\nNO", "non-prime squares"

# large numbers
assert run("3\n999900000000 1000000000000 1000003*1000003\n") == "NO\nNO\nYES", "large numbers"
Test input Expected output What it validates
1 NO minimum input
4 9 25 49 YES YES YES YES correct identification of T-primes
16 36 100 NO NO NO squares of non-primes rejected
999900000000 1000000000000 1000003*1000003 NO NO YES handles large numbers correctly

Edge Cases

The algorithm handles 1 correctly because 1 is not prime and thus not a T-prime. Large numbers like 10^12 are handled accurately because we avoid floating-point square roots and rely on integer multiplication. Numbers that are perfect squares but whose roots are composite, like 16, 36, and 100, are rejected because the square root is checked against the precomputed prime set. The set lookup guarantees constant-time verification even at the upper bound of 100,000 numbers.