[백준] 9421 소수상근수 python

https://www.acmicpc.net/problem/9421

9421번: 소수상근수

소수인 상근수를 찾는 문제이다.

소수를 에라토스테네스의 체로 구해놓고 소수가 상근수인지 찾으면 된다.

from math import isqrt
def findPrime(N):
    prime = [1 for _ in range(N + 1)]
    prime[0] = prime[1] = 0
    for i in range(2, isqrt(N) + 1):
        if prime[i]:
            for j in range(i + i, N + 1, i):
                prime[j] = 0
    return prime

def sanggn(n):
    visit = [0 for _ in range(100*7)]
    tmp = 0
    while 1:
        while n:
            tmp += (n % 10) ** 2
            n //= 10

        if tmp == 1:
            return True

        if visit[tmp]:
            return False
        visit[tmp] = 1
        n = tmp
        tmp = 0

def main():
    n = int(input())
    prime = findPrime(n)
    answer = []
    for i in range(2, n + 1):
        if prime[i] and sanggn(i):
            answer.append(i)

    print('\n'.join(map(str, answer)))


if __name__ == "__main__":
    main()

 

여기에서 visit을 함수 내에서 계속 초기화하면서 사용했는데 얘를 글로벌 변수로 사용하고 해당 수가 상근 수인지 아닌지 DP 형식으로 하면 시간을 더 줄일수 있을것같다.