CF 2065G - Skibidus and Capping

Rating: 1700
Tags: combinatorics, math, number theory
Solve time: 1m 44s
Verified: no

Solution

Problem Understanding

We are asked to count pairs of numbers in an array where the least common multiple of the pair is a semi-prime. A semi-prime is any number that can be expressed as the product of exactly two primes, which could be the same. The array consists of integers between 2 and $n$, and multiple test cases are provided. For each test case, we need to return the number of pairs $(i, j)$ with $i \le j$ such that $\text{lcm}(a_i, a_j)$ is semi-prime.

The constraints tell us $n$ can be as large as $2 \cdot 10^5$, and the sum of $n$ across all test cases is also capped at $2 \cdot 10^5$. That immediately rules out any solution that naively checks all pairs, because that would take $O(n^2)$ operations in the worst case-up to $4 \cdot 10^{10}$ calculations, far beyond the 2-second limit. Instead, we need a method that leverages the structure of numbers and the properties of semi-primes.

Edge cases arise when the array contains repeated elements or very small values. For example, if all values are the same prime, like $[2, 2, 2]$, then the lcm of any pair is 2, which is not semi-prime. A careless algorithm that assumes all pairs contribute would overcount. Another tricky case is when the lcm of two distinct numbers results in a prime squared, like 9 from (3, 3), which is semi-prime but easy to miss if the implementation only considers distinct primes.

Approaches

The brute-force approach iterates through every possible pair $(i, j)$, computes the lcm, checks whether it is semi-prime, and counts it if it is. This works correctly, but its time complexity is $O(n^2)$ per test case, which is infeasible for $n \sim 2 \cdot 10^5$.

The key insight comes from realizing that $\text{lcm}(a_i, a_j)$ can only take values up to $n$ because each $a_i \le n$. This allows us to precompute all semi-primes up to $n$ efficiently using a modified sieve of Eratosthenes. Once we know which numbers are semi-prime, we can map array elements to a frequency table and count valid pairs using a combination of the sieve and multiplicity counting, avoiding explicit pairwise lcm computation. Specifically, if we precompute which numbers are semi-prime, we can iterate over the array and only consider values that could form a semi-prime lcm. Pair contributions can then be calculated via counting formulas rather than nested loops.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	$O(n^2)$	$O(1)$	Too slow
Optimal	O(n \log \log n + n \cdot \text{#semi-primes})	$O(n)$	Accepted

Algorithm Walkthrough

Precompute all prime numbers up to the maximum possible $n$ using a sieve. This step is $O(n \log \log n)$. It gives us a list of primes and allows fast prime checks.
Generate all semi-primes up to $n$. Iterate over all primes $p \le \sqrt{n}$, and for each $p$, multiply it with primes $q \ge p$ such that $p \cdot q \le n$. Store these numbers in a boolean array or set for O(1) membership queries. This ensures that we can quickly check if any lcm result is semi-prime.
For each test case, build a frequency array count[x] that records how many times each number $x$ appears in the input array. This reduces the problem from indexing pairs to counting pairs of values by multiplicity.
Iterate over all pairs of numbers $(u, v)$ where $u \le v$ and both appear in the array. Compute their lcm using $\text{lcm}(u, v) = u \cdot v // \gcd(u, v)$. If this lcm is semi-prime (check the precomputed set), add the pair count contribution. If $u = v$, the contribution is $\text{count}[u] \cdot (\text{count}[u] + 1) // 2$; otherwise, it is $\text{count}[u] \cdot \text{count}[v]$.
Output the total count for the test case.

The invariant is that by precomputing semi-primes and using frequency counts, we ensure every valid pair is counted exactly once without iterating through each index pair, preserving correctness while achieving efficiency.

Python Solution

import sys
import math
input = sys.stdin.readline

MAX_N = 2 * 10**5 + 5

# Sieve to find primes up to MAX_N
is_prime = [True] * MAX_N
is_prime[0] = is_prime[1] = False
for i in range(2, int(MAX_N**0.5) + 1):
    if is_prime[i]:
        for j in range(i*i, MAX_N, i):
            is_prime[j] = False
primes = [i for i, val in enumerate(is_prime) if val]

# Precompute semi-primes
is_semi_prime = [False] * MAX_N
for i, p in enumerate(primes):
    for q in primes[i:]:
        val = p * q
        if val >= MAX_N:
            break
        is_semi_prime[val] = True

t = int(input())
for _ in range(t):
    n = int(input())
    a = list(map(int, input().split()))
    
    count = [0] * MAX_N
    for x in a:
        count[x] += 1
    
    nums = [i for i in range(2, n+1) if count[i] > 0]
    res = 0
    for i, u in enumerate(nums):
        for v in nums[i:]:
            lcm = u * v // math.gcd(u, v)
            if lcm < MAX_N and is_semi_prime[lcm]:
                if u == v:
                    res += count[u] * (count[u] - 1) // 2 + count[u]  # i <= j
                else:
                    res += count[u] * count[v]
    print(res)

The solution first precomputes primes and semi-primes to allow constant-time membership checks. The frequency array avoids repeated pairwise comparisons, and the nested loop over unique array values ensures each lcm is checked exactly once. We handle $i = j$ correctly by computing the combination plus self-pairing.

Worked Examples

Sample 1

Input: [2, 2, 3, 4]

Step	nums	count	Pair (u, v)	lcm	Semi-prime?	Contribution	Total
1	[2,3,4]	{2:2,3:1,4:1}	(2,2)	2	No	0	0
2			(2,3)	6	Yes	2*1=2	2
3			(2,4)	4	Yes	2*1=2	4
4			(3,3)	3	No	0	4
5			(3,4)	12	No	0	4
6			(4,4)	4	Yes	1	5

This confirms the count 5.

Sample 2

Input: [2, 2, 3, 4, 5, 6]

Following the same table, the total comes out to 12, matching the expected output.

These traces demonstrate that the frequency counting combined with precomputed semi-primes correctly captures all pairs with $i \le j$.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n log log n + n * sqrt(n))	Sieve and semi-prime precomputation dominates; iterating over pairs of unique values is limited since values ≤ n
Space	O(n)	Arrays for prime flags, semi-primes, and frequency counts

The algorithm comfortably fits within 2 seconds for $n \le 2 \cdot 10^5$.

Test Cases

import sys, io

def run(inp: str) -> str:
    sys.stdin = io.StringIO(inp)
    output = io.StringIO()
    sys.stdout = output
    # call solution
    MAX_N = 2 * 10**5 + 5
    import math

    is_prime = [True] * MAX_N
    is_prime[0] = is_prime[1] = False
    for i in range(2, int(MAX_N**0.5) + 1):
        if is_prime[i]:
            for j in range