[파이썬, Python] 백준 9421번 : 소수상근수

백준 9421번 : 소수상근수

(문제 바로가기)

 

 

내 코드

from collections import defaultdict

def get_prime(num):  # 1부터 num까지의 범위에서 소수 찾아서 리스트로 반환
    prime = [False, False] + [True] * (num - 1)
    for i in range(2, int(num ** 0.5) + 1):
        if not prime[i]:
            continue
        j = 2
        while i * j <= num:
            prime[i * j] = False
            j += 1
    return [i for i in range(2, num + 1) if prime[i]]


n = int(input())
answer = set()
visited = defaultdict(set)  # 방문
for p in get_prime(n):
    temp = str(p)  # 각 자리수 접근을 위해 타입 변경
    while 1:
        temp = sum(map(lambda x: int(x) ** 2, temp))  # 각 자리수의 제곱의 합
        if temp == 1 or (len(visited[temp]) >= 1 and answer & visited[temp]):  # 제곱의 합이 1이거나 이 합이 소수상근수가 되는 결과라면
            answer.add(p)
            visited[temp].add(p)
            break
        elif len(visited[temp]) >= 1 and not (answer & visited[temp]):  # 이 합의 결과가 소수상근수가 될 수 없는 수라면
            break
        else:
            visited[temp].add(p)
        temp = str(temp)

print(*sorted(answer), sep='\n')