题解 | 小红的“质数”寻找

# read from stdin
import sys

lines = []

for line in sys.stdin:
    lines.append(line.rstrip("\n"))

total_xs = int(lines.pop(0))


def find(x):
    """
    # first digit less than 5
        find a number that satisfy the sum is prime
        the easiest:
            if it is a prime, return it and 00000 suffix
            if not, grow + 1 to find a digit
    # first digit gt 5
        find a number that satisfy the sum - 1 is prime
        the easiest: 10000001 where 1 is the additional one

    """
    dict_num_to_nearest_prime = {
        1: 2,
        2: 3,
        3: 5,
        4: 5,
        5: 7,
        6: 7,
        # 8: 1,
        # 9: 1,
    }
    # if it is single-digited, might be -1
    if len(x) == 1:
        num = int(x[0])
        if num == 0: 
            return -1
        if num in dict_num_to_nearest_prime: 
            return int(dict_num_to_nearest_prime[num])

    digits = list(map(int, list(x)))

    if digits[0] <= 5:
        first_digit = str(dict_num_to_nearest_prime[digits[0]])
        rest = "0" * (len(digits) - 1) 
        return int(first_digit + rest)
        
    else:
        additional_digit = '1'
        rest =  '0' * (len(digits) - 1) + '1'
        return int(additional_digit + rest)
    
def is_prime(i):
    if i <= 1:
        return False
    for n in range(2, int(i**0.5) + 1):
        if i % n == 0:
            return False
    return True


for _ in range(total_xs):
    x = lines.pop(0)
    result = find(x)
    assert result >= int(x)
    assert result <= 2*int(x)
    assert is_prime(sum(map(int, str(result))))
    print(result)