CF 1033B - Square Difference

Rating: 1100
Tags: math, number theory
Solve time: 4m 51s
Verified: yes

Solution

Problem Understanding

We start with a large square piece of cloth with side length $a$. From one corner, a smaller square of side $b$ is cut out. The remaining cloth is an L-shaped region whose area is simply the area of the big square minus the area of the removed square, so $a^2 - b^2$.

The task is not to compute this area directly for output, but to decide whether this resulting number is a prime.

Each test case gives a pair $(a, b)$, and we must answer whether $a^2 - b^2$ is prime.

The constraints are very large: $a$ and $b$ can go up to $10^{11}$, which means $a^2$ can reach $10^{22}$. That immediately rules out any approach that tries to factor numbers by trial division up to the value itself or even up to its square root in a naive way for every test case. Even $O(\sqrt{n})$ per test case becomes borderline if done repeatedly, because $\sqrt{10^{22}} = 10^{11}$, which is far too large.

The key structure here is that we are not given an arbitrary number. The value is always a difference of two squares, which has strong algebraic properties that restrict its factorization behavior.

A subtle edge case appears when $a$ and $b$ are very close. For example, $a=34, b=33$ gives $34^2 - 33^2 = 67$, which is prime. But for slightly larger gaps like $16, 13$, we get $87$, which is composite even though the numbers are small. Another edge case is when the expression becomes even; any even number greater than 2 is automatically non-prime, and this already rules out many cases without computation.

Approaches

A brute-force approach would compute $a^2 - b^2$ and then test primality using trial division up to $\sqrt{a^2 - b^2}$. While correct, this is infeasible because the difference can be as large as $10^{22}$, making primality testing require up to $10^{11}$ iterations per test case in the worst case.

The key insight comes from rewriting the expression:

$$a^2 - b^2 = (a-b)(a+b)$$

This factorization is decisive. The number is always a product of two integers greater than zero, and its structure is extremely constrained.

Since $a > b$, both factors are positive integers. Also, $a+b > a-b$, and both are at least 1. The only way a product of two integers is prime is when one of them equals 1.

So we ask: when can $(a-b)(a+b)$ equal a prime?

A prime number has exactly two positive divisors: 1 and itself. Therefore one of the factors must be 1. Since $a+b \ge 2$, it cannot be 1. So the only possibility is:

$$a - b = 1$$

If $a-b = 1$, then the expression becomes:

$$(1)(2b+1) = 2b+1$$

So the result is prime if and only if $2b+1$ is prime.

This reduces the problem to a single primality test on a number up to $2 \cdot 10^{11}$, which is manageable with deterministic trial division up to $\sqrt{n}$, or a faster check using only odd divisors.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	$O(\sqrt{n} \cdot t)$ with $n \approx 10^{22}$	$O(1)$	Too slow
Optimal	$O(\sqrt{b})$ per valid case (often constant due to early rejection)	$O(1)$	Accepted

Algorithm Walkthrough

Key idea setup

Compute $d = a - b$. This isolates the only part that can possibly reduce the factorization into something prime-like.
If $d \neq 1$, immediately return "NO". This follows because the number becomes a product of two integers both strictly greater than 1, so it is composite.
If $d = 1$, compute $x = a^2 - b^2 = (a-b)(a+b) = 2b + 1$.
Check whether $x$ is prime using trial division up to $\sqrt{x}$.
Return "YES" if prime, otherwise "NO".

The crucial observation is that all complexity collapses into a single special case where the structure degenerates into a candidate prime.

Why it works

The invariant is that every valid value $a^2 - b^2$ factors exactly as $(a-b)(a+b)$. Since $a+b > a-b$ and both are integers greater than or equal to 1, the expression is composite unless one factor is 1. The only possible way to satisfy this is $a-b = 1$, which uniquely reduces the problem to testing primality of a single derived number $2b+1$. No other structure can produce a prime because any nontrivial factorization immediately violates primality.

Python Solution

import sys
input = sys.stdin.readline

def is_prime(n: int) -> bool:
    if n < 2:
        return False
    if n % 2 == 0:
        return n == 2
    i = 3
    while i * i <= n:
        if n % i == 0:
            return False
        i += 2
    return True

t = int(input())
for _ in range(t):
    a, b = map(int, input().split())
    
    if a - b != 1:
        print("NO")
    else:
        val = 2 * b + 1
        print("YES" if is_prime(val) else "NO")

The implementation relies on separating the trivial rejection case $a-b \neq 1$ from the only meaningful case. This avoids any large number arithmetic on $a^2$ directly, preventing overflow concerns and unnecessary computation.

The primality check is optimized by skipping even numbers after handling 2, which is sufficient given the structure of $2b+1$, always odd.

Worked Examples

Example 1: $a=16, b=13$

We compute $d = a-b = 3$, so we immediately reject.

Step	a	b	a-b	Decision
Init	16	13	-	-
Check diff	16	13	3	Reject

This demonstrates that any difference greater than 1 leads to a guaranteed composite number.

Example 2: $a=34, b=33$

Now $d = 1$, so we proceed.

We compute $x = 2b+1 = 67$, then test primality.

Step	a	b	a-b	value (2b+1)	Prime check
Init	34	33	-	-	-
Diff	34	33	1	-	proceed
Compute	34	33	1	67	test
Result	34	33	1	67	prime

This confirms the special case where the structure reduces to a prime candidate.

Complexity Analysis

Measure	Complexity	Explanation
Time	$O(t \sqrt{b})$ worst case	Only when $a-b=1$, otherwise constant time rejection
Space	$O(1)$	Only a few integers used

The constraints allow up to 5 test cases, so even the worst-case primality check on a 10-digit number is easily fast enough.

Test Cases

import sys, io

def run(inp: str) -> str:
    sys.stdin = io.StringIO(inp)
    input = sys.stdin.readline

    def is_prime(n: int) -> bool:
        if n < 2:
            return False
        if n % 2 == 0:
            return n == 2
        i = 3
        while i * i <= n:
            if n % i == 0:
                return False
            i += 2
        return True

    t = int(input())
    out = []
    for _ in range(t):
        a, b = map(int, input().split())
        if a - b != 1:
            out.append("NO")
        else:
            out.append("YES" if is_prime(2 * b + 1) else "NO")
    return "\n".join(out)

# provided samples
assert run("4\n6 5\n16 13\n61690850361 24777622630\n34 33\n") == "YES\nNO\nNO\nYES"

# custom cases
assert run("1\n3 2\n") == "YES"
assert run("1\n10 9\n") == "NO"
assert run("1\n5 4\n") == "YES"
assert run("1\n100 98\n") == "NO"

Test input	Expected output	What it validates
3 2	YES	smallest valid difference case
10 9	NO	composite odd result
5 4	YES	small prime result case
100 98	NO	difference not equal to 1

Edge Cases

Case: $a-b > 1$

For input $a=10, b=7$, we get $a-b=3$. The algorithm immediately returns "NO". This is correct because $a^2-b^2 = (a-b)(a+b) = 3 \cdot 17 = 51$, which is composite.

The computation never reaches primality testing, which avoids unnecessary work.

Case: minimal difference

For $a=3, b=2$, we get $a-b=1$, so we compute $2b+1=5$. The primality test confirms 5 is prime, so output is "YES". This is the smallest non-trivial case where the logic activates the second branch.

Case: large values

For $a=10^{11}, b=10^{11}-1$, the algorithm computes $2b+1 \approx 2 \cdot 10^{11}$. The primality check runs in $O(\sqrt{b})$, but this still fits easily because only a few hundred thousand iterations are needed, and $t \le 5$.

This confirms the solution scales cleanly even at the upper bound.