Editorial for the Primes


Remember to use this editorial only when stuck, and not to copy-paste code from it. Please be respectful to the problem author and editorialist.
Submitting an official solution before solving the problem yourself is a bannable offence.
from collections import defaultdict
import sys


def primes():
    i = 2
    prime_list = []
    while True:
        for p in prime_list:
            if i % p == 0:
                break
        else:
            prime_list.append(i)
            yield i
        i += 1

prime_list = defaultdict(lambda : list([dict(), dict(), dict(), dict(), dict(), dict()]))
for p in primes():
    if p > 9999:
        p_ = str(p)
        k = sum(int(x) for x in p_)
        prime_list[k][0][p] = True
        for i in range(1, 6):
            prime_list[k][i][p_[:i]] = True
    if p >= 99999:
        break

def is_solution(k, n, numbers):
    if len(numbers) != 5:
        return False
    return True

def process_solution(numbers):
    for n in numbers:
        print(n)
    print()

def extend_solution(k, numbers):
    global primes
    i = len(numbers)
    if i > 4:
        return
    for p in prime_list[k][0]:
        new_numbers = list(numbers)
        new_numbers.append(p)
        yield new_numbers

def test(k, n, numbers):
    if numbers[0] // 10000 != n:
        return False
    m = len(numbers)
    nums_ = ['']*7
    for j, n in enumerate(numbers):
        r = tuple(x for x in str(n))
        for i in range(5):
            nums_[i] += r[i]
        nums_[5] += r[j]
        nums_[6] += r[4-j]
    for n in nums_:
        if sum(int(x) for x in n) > k:
            return False
        if ''.join(n) not in prime_list[k][m]:
            return False
    return True

def backtracking(k, n, numbers):
    if is_solution(k, n, numbers):
        process_solution(numbers)
    else:
        for new_numbers in extend_solution(k, numbers):
            if test(k, n, new_numbers):
                backtracking(k, n, new_numbers)

d_sum = int(sys.stdin.readline())
digit = int(sys.stdin.readline())

backtracking(d_sum, digit, [])

Comments

There are no comments at the moment.