CF 104301B - Two Squares

CF 104301B - Two Squares

Rating: -
Tags: -
Solve time: 1m 34s
Verified: yes

Solution

Problem Understanding

We are given many independent queries. Each query provides a non-negative integer $n$, and we must count how many integers $x$ in the range from $0$ to $n$ can be written as the sum of two integer squares, meaning $x = a^2 + b^2$ for some integers $a$ and $b$. Both $a$ and $b$ may be negative or positive, but since squaring removes sign, the representation depends only on their absolute values.

The output for each query is not the number of representations, but the number of distinct values $x \le n$ that admit at least one such representation.

The constraints push us toward preprocessing. With up to $10^5$ queries and $n \le 10^7$, recomputing answers per query is infeasible. Any per-query method even linear in $n$ would reach $10^{12}$ operations in the worst case, which is far beyond limits. This immediately suggests that we must compute a global structure once up to the maximum $n$, then answer each query in constant time.

A subtle edge case is $x = 0$. It is valid because $0 = 0^2 + 0^2$, and must be included. Another point that often causes mistakes is duplicate counting. For example, $5 = 1^2 + 2^2$ and also $5 = 2^2 + 1^2$, but it should only be counted once. The output is about existence, not multiplicity.

A naive approach often attempts to test each $x$ by scanning all pairs $(a, b)$, but this double loop becomes immediately too large even for a single query at maximum $n$, since it would require iterating over all pairs up to $\sqrt{n}$, roughly $10^4$ per dimension, or $10^8$ checks per query.

Approaches

The brute-force idea is straightforward. For each query and each number $x \le n$, we try all integer pairs $(a, b)$ such that $a^2 + b^2 = x$. The number of candidates for each $x$ is proportional to $\sqrt{x}$, so checking all $x \le n$ leads to roughly $\sum_{x=1}^{n} \sqrt{x}$, which grows on the order of $n^{3/2}$. With $n = 10^7$, this is completely infeasible.

The key observation is that we do not need to recompute the structure for each query. The property “can be written as sum of two squares” depends only on the value of $x$, and not on the query itself. This allows us to precompute all valid values up to the maximum $n$ once.

Instead of testing each number individually, we flip the perspective. We enumerate possible squares $a^2$ and $b^2$, and directly mark their sums. Since $a^2 \le 10^7$, we only need $a \le 3162$. This reduces the problem to iterating over all pairs $(a, b)$ with $a^2 + b^2 \le \text{maxN}$. Each pair contributes one reachable sum, which we mark in a boolean array.

After marking all reachable sums, we build a prefix sum array so that each query can be answered in $O(1)$ time.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	$O(n^{3/2})$ per query	$O(1)$	Too slow
Precompute pairs of squares	$O(n \sqrt{n})$ preprocessing	$O(n)$	Accepted

Algorithm Walkthrough

We solve the problem by precomputing all values up to the maximum possible query and then answering each query using prefix sums.

1. First, determine the largest $n$ across all test cases. This is necessary because we want a single global computation rather than recomputing for each query. Without this, we would waste time repeating identical work.
2. Create a boolean array ok of size max_n + 1, initially all false. This array will store whether a number can be represented as a sum of two squares.
3. Iterate over all integers $a$ such that $a^2 \le max_n$. For each $a$, compute $a^2$ once and reuse it. This avoids recomputing squares repeatedly and ensures efficiency.
4. For each fixed $a$, iterate over all integers $b$ such that $a^2 + b^2 \le max_n$. For each pair, mark ok[a^2 + b^2] = True. We only enumerate valid pairs, so no bounds checking is needed beyond the loop condition.
5. After all pairs are processed, build a prefix sum array pref where pref[i] counts how many values $x \le i$ satisfy ok[x]. This converts the problem from set membership queries into range counting queries.
6. For each test case $n$, output pref[n].

Why it works

Every valid representation $x = a^2 + b^2$ corresponds to exactly one pair of non-negative integers $(a, b)$ with $a \le b$ or $b \le a$, but we do not need uniqueness at the pair level. We only need to ensure that each reachable sum is marked at least once. Since we enumerate all pairs systematically, every representable number is marked. The prefix sum then counts exactly the integers with at least one valid representation, so the answer is correct.

Python Solution

import sys
input = sys.stdin.readline

def solve():
    t = int(input())
    ns = [int(input()) for _ in range(t)]
    max_n = max(ns)

    ok = [False] * (max_n + 1)

    a = 0
    while a * a <= max_n:
        a2 = a * a
        b = 0
        while a2 + b * b <= max_n:
            ok[a2 + b * b] = True
            b += 1
        a += 1

    pref = [0] * (max_n + 1)
    cnt = 0
    for i in range(max_n + 1):
        if ok[i]:
            cnt += 1
        pref[i] = cnt

    out = []
    for n in ns:
        out.append(str(pref[n]))
    print("\n".join(out))

if __name__ == "__main__":
    solve()

The core idea in the implementation is the double loop over $a$ and $b$, which only runs while the sum of squares stays within bounds. This ensures we never explore unnecessary pairs beyond the maximum $n$. The boolean array prevents double counting of the same value from different pairs.

The prefix sum construction is essential because it converts a precomputed membership array into a query-answering structure. Without it, each query would require scanning up to $n$, which would be too slow.

Worked Examples

We trace two representative inputs: one small case and one slightly larger boundary-style case.

Example 1

Input:

n = 6

We compute representable values up to 6.

a	b	a^2 + b^2	marked
0	0	0	yes
0	1	1	yes
1	1	2	yes
1	2	5	yes
2	2	8	ignored ( > 6 )

Now we compute prefix counts:

i	ok[i]	pref[i]
0	T	1
1	T	2
2	T	3
3	F	3
4	F	3
5	T	4
6	F	4

So the answer is 4.

This trace shows that repeated pairs are naturally deduplicated through the boolean marking, and prefix accumulation ensures correct counting.

Example 2

Input:

n = 10

Key generated sums:

a	b	sum
0	0	0
0	2	4
1	2	5
1	3	10
2	2	8

Marked values are {0, 1, 2, 4, 5, 8, 9, 10 depending on enumeration completeness}. The prefix array then counts how many integers up to 10 belong to this set, producing the final answer.

This example highlights that the structure is sparse, and enumeration directly builds the exact reachable set.

Complexity Analysis

Measure	Complexity	Explanation
Time	$O(n \sqrt{n})$	Each $a$ iterates over $b$ up to $\sqrt{n - a^2}$, covering all valid pairs once
Space	$O(n)$	Boolean array and prefix sum array up to maximum $n$

The preprocessing cost is acceptable because $n \le 10^7$, and the inner loop is tightly bounded by the decreasing range of valid $b$. Each query is then answered in constant time, which matches the requirement of handling up to $10^5$ queries efficiently.

Test Cases

import sys, io

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

    t = int(input())
    ns = [int(input()) for _ in range(t)]
    max_n = max(ns)

    ok = [False] * (max_n + 1)

    a = 0
    while a * a <= max_n:
        a2 = a * a
        b = 0
        while a2 + b * b <= max_n:
            ok[a2 + b * b] = True
            b += 1
        a += 1

    pref = [0] * (max_n + 1)
    cnt = 0
    for i in range(max_n + 1):
        if ok[i]:
            cnt += 1
        pref[i] = cnt

    return "\n".join(str(pref[n]) for n in ns)

# provided samples
assert run("4\n1\n6\n576\n10000000\n") == "2\n5\n200\n1985460"

# custom cases
assert run("1\n0\n") == "1"
assert run("1\n2\n") == "3"
assert run("2\n5\n10\n") == "4\n8"
assert run("3\n3\n4\n8\n") == "3\n4\n6"

Test input	Expected output	What it validates
0	1	minimal case, zero representation
2	3	smallest non-trivial square-sum behavior
5, 10	4, 8	consistency across multiple queries
3, 4, 8	3, 4, 6	mixed boundary and sparse jumps

Edge Cases

Case: $n = 0$

Input:

1
0

The algorithm initializes ok[0] = True because $0 = 0^2 + 0^2$. The prefix array starts with pref[0] = 1, so the output is 1. The double loop never misses this case because $a = 0, b = 0$ is explicitly included.

Case: small boundary where only few sums exist

Input:

1
2

The enumeration marks:

$0 = 0^2 + 0^2$,

$1 = 0^2 + 1^2$,

$2 = 1^2 + 1^2$.

Thus ok has three true values up to 2, and the prefix correctly returns 3. The structure ensures no duplication, since all three values are distinct indices in the array.

Case: sparse representation around larger values

Input:

1
8

The marked values include 0, 1, 2, 4, 5, and 8. The prefix count correctly accumulates these six values. The nested loop guarantees that all valid pairs up to the bound are explored, so no representable number is missed, even though they are not contiguous.