Magic square of squares

"""Verifier for "Magic square of squares" (statement in magic_square_of_squares.json).

Problem: find a 3x3 magic square whose nine entries are distinct perfect
squares. The solver submits the nine square roots r[0..8] in row-major order,
so the grid is

    r[0]**2  r[1]**2  r[2]**2
    r[3]**2  r[4]**2  r[5]**2
    r[6]**2  r[7]**2  r[8]**2

and every row, every column and both diagonals must have the same sum.

Expected submission: a JSON object with one key, a list of 9 decimal strings:

    {"roots": ["1", "2", "3", "4", "5", "6", "7", "8", "9"]}

What the verifier checks, in order:
  1. The payload is at most MAX_BYTES characters and is a JSON object with
     exactly the key "roots".
  2. "roots" is a list of exactly 9 well-formed decimal strings of at most
     MAX_DIGITS digits each (see _parse_int).
  3. Every root is at least 1.
  4. The nine cells (the squared roots) are pairwise distinct. Without this,
     nine copies of the same square would trivially be "magic".
  5. All 8 lines (3 rows, 3 columns, 2 diagonals) have the same sum.
"""
import json

# When True (the real problem), each cell is the square of the submitted root.
# The tests set it to False, so the submitted numbers are used as the cells
# directly, to check that an ordinary magic square (the Lo Shu) is accepted.
REQUIRE_SQUARES = True
# Maximum number of digits allowed in each root.
MAX_DIGITS = 10_000
# Maximum size of the whole submission, in characters.
MAX_BYTES = 100_000

# The 8 lines of the grid, as triples of cell indices (row-major, 0 = top left):
# three rows, then three columns, then the two diagonals.
LINES = ((0, 1, 2), (3, 4, 5), (6, 7, 8), (0, 3, 6), (1, 4, 7), (2, 5, 8), (0, 4, 8), (2, 4, 6))


# ---------------------------------------------------------------------------
# Input-parsing helpers. The same code is copied into every verifier so that
# each file works on its own.
# ---------------------------------------------------------------------------

def _parse_int(s, max_digits, allow_neg=False):
    """Convert a strictly formatted decimal string to an int, or raise ValueError.

    Accepted: "0", "42", and "-42" when allow_neg is True.
    Rejected: anything that is not a string (for example a JSON number), "",
    "+42", " 42", "007", "-0", non-ASCII digits such as "٣", and strings with
    more than max_digits digits.
    """
    # Cheap length check before any real work: at most max_digits digits,
    # plus one character for an optional minus sign.
    if not isinstance(s, str) or len(s) > max_digits + 1:
        raise ValueError
    # Split off the sign. A "-" is only allowed where negative values make sense.
    neg = s.startswith("-")
    if neg and not allow_neg:
        raise ValueError
    body = s[1:] if neg else s
    # What remains must be 1 to max_digits ASCII digits with no leading zero
    # ("0" on its own is fine). isascii() is needed because isdigit() also
    # accepts digits from other scripts.
    if not (body.isascii() and body.isdigit()) or len(body) > max_digits or (len(body) > 1 and body[0] == "0"):
        raise ValueError
    # "-0" is not a canonical way to write zero.
    if neg and body == "0":
        raise ValueError
    # Python (3.11 and later) refuses int() on strings longer than 4300 digits
    # by default, so the number is built from chunks of 4000 digits.
    n = 0
    for i in range(0, len(body), 4000):
        chunk = body[i:i + 4000]
        n = n * 10 ** len(chunk) + int(chunk)
    return -n if neg else n


def _load(solution, keys, max_bytes):
    """Parse the submission as a JSON object with exactly the given keys, or raise."""
    # Refuse oversized payloads before handing them to the JSON parser.
    if not isinstance(solution, str) or len(solution) > max_bytes:
        raise ValueError
    d = json.loads(solution)
    # It must be an object (not a list, string or number) whose keys are
    # exactly the expected ones: nothing missing, nothing extra.
    if not isinstance(d, dict) or set(d) != set(keys):
        raise ValueError
    return d


# ---------------------------------------------------------------------------
# The actual check.
# ---------------------------------------------------------------------------

def _check(solution):
    # Step 1: read the JSON object with its single key "roots".
    sol = _load(solution, ("roots",), MAX_BYTES)

    # Step 2: "roots" must be a list of exactly nine numbers.
    roots = sol["roots"]
    if not isinstance(roots, list) or len(roots) != 9:
        return False
    rs = [_parse_int(r, MAX_DIGITS) for r in roots]

    # Step 3: every root must be positive.
    if min(rs) < 1:
        return False

    # Build the grid: square each root (or use it directly in test mode).
    cells = [r * r for r in rs] if REQUIRE_SQUARES else rs

    # Step 4: all nine cells must be different from each other.
    if len(set(cells)) != 9:
        return False

    # Step 5: the 8 lines must all have the same sum, i.e. their sums must
    # contain exactly one distinct value.
    line_sums = {sum(cells[i] for i in line) for line in LINES}
    return len(line_sums) == 1


def verify(solution: str, rng) -> bool:
    """Platform entry point: return True only for a valid, correct solution.

    `rng` (a numpy.random.Generator) is part of the platform interface but is
    not used: this check is deterministic. Any exception (malformed JSON, a bad
    number, ...) means the submission is rejected, so this function never raises.
    """
    try:
        return _check(solution) is True
    except Exception:
        return False