[알고리즘] 백준 10434 파이썬 - 행복한 소수

 

10434번: 행복한 소수

문제

Q. : 아래의 수열에서 다음에 올 수를 찾으시오.

313 331 367 ...

경복 : ??

강산 : 379.

경복 : 뭐?

강산 : 행복한 소수잖아.

경복 : 행복한 뭐?

강산 : 그러니까, 자리수의 제곱의 합을 구하는 연산을 계속 반복했을 때 1이 되는 수를 행복한 수라고 하잖아. 행복한 소수는 그 중 소수인 수이고.

7은 분명 소수이다. 과연 행복할까?

• 7 → 72 = 49
• 49 → 42 + 92 = 97
• 97 → 92 + 72 = 130
• 130 → 12 + 32 + 02 = 10
• 10 → 12 + 02 = 1

7은 행복한 수이다 ☺.

사실 7은 행복한 소수 중 가장 작은 수이다. (이 문제에서는 1을 소수가 아닌 것으로 생각한다)

어떤 수가 주어지면 이 수가 행복한 소수인지 판정해보자.

문제풀이

import sys
input = sys.stdin.readline

#01
def happy(n):
    arr=[]
    x=0
    while True:
        if n<10:
            a1=n
            x=int(a1**2)
        elif n<100:
            a2=int(n/10)
            a1=n%10
            x=int(a2**2+a1**2)
        elif n<1000:
            a3=int(n/100)
            a2=int((n%100)/10)
            a1=n%10
            x=int(a3**2+a2**2+a1**2)
        elif x<10000:
            a4=int(n/1000)
            a3=int((n%1000)/100)
            a2=int((n%100)/10)
            a1=n%10
            x=int(a4**2+a3**2+a2**2+a1**2)
        if x==1:
            return 1
        if x in arr:
            return 0
        arr.append(x)
        n=x
#02  
n=10000
tmp=[False,False]+([True]*(n-1))
for i in range(2,int(n**0.5)+1):
    for j in range(2*i,n+1,i):
        if tmp[j]:
            tmp[j]=False
for _ in range(int(input())):
    a,b=map(int,input().split())
    #03
    if tmp[b]:
        x=happy(b)
        if x==1:
            print(a ,b ,"YES")
        else:
            print(a ,b ,"NO")
    else:
        print(a ,b ,"NO")

- #01 : 행복한 소수인지를 판별하는 함수이다.

           처음에 반복문 코드로 간단하게 구현했을 때 시간초과가 발생하였다.

           그래서 자릿수마다 계산하여 행복한 소수가 되는지 확인하였다.

           1이면 그대로 return 1을 하고, 그렇지 않으면 arr에 추가한다.

           loop 즉 arr에 있는 x 값이라면 return 0을 해준다.

 

- #02 : 에라토스테네스의 체를 이용하여 10000까지의 소수를 구해준다.

- #03 : b값이 소수라면 happy 함수를 실행한 결과값에 따라 출력해주고,

           b값이 소수가 아니라면 바로 NO를 출력해준다.