[PS] BOJ 10434 / 행복한 소수

문제 링크: https://www.acmicpc.net/problem/10434

Thumbnail: Photo by The Travel Nook / Unsplash

'행복한 수' 를 판별하는 과정이 바로 떠오르지 않았지만, 패턴을 고찰해 보니 쉽게 이해할 수 있었습니다!

풀이

임의의 수 \(M\)가 행복한 소수인지 판단하려면, 다음 두 가지 과정을 거치면 된다.

• \(M\)이 소수인가?
• \(M\)이 행복한 수인가?

소수인지 판별하는 경우는 문제에서 \(1 \le M \le 10,000\) 라는 조건을 주었으니, 그냥 이 범위 안의 소수를 에라토스테네스의 체 알고리즘으로 미리 계산해 두면 된다.

행복한 수인지 판별하는 경우를 생각해 보자. 행복한 수란, 각 자리의 제곱의 합 계산하는 과정을 반복해 그 결과가 1이 되는 수를 말한다.

대표적인 예시로, 가장 작은 행복한 소수인 7이 행복한 수인지 판단하는 과정을 보자.

• 7 → 7² = 49
• 49 → 4² + 9² = 97
• 97 → 9² + 7² = 130
• 130 → 1² + 3² + 0² = 10
• 10 → 1² + 0² = 1

각 자리의 제곱의 합을 계산하고, 그 결과가 1인지 비교하는 과정을 먼저 구현해보자.

while (condition):
    sum = 0
    while n > 0:
        sum += (n % 10) ** 2
        n //= 10
    if sum == 1:
        return True
    n = sum

하지만, 이대로는 행복한 수가 아닌 경우에는 while문을 벗어날 수 없다! 이를 조절할 적절한 조건을 찾아야 한다. 행복한 수가 아닌 어떤 수에 대해, 그 계산 과정을 분석해 보자.

• 8 → 8² = 64
• 64 → 6² + 4² = 52
• 52 → 5² + 2² - 29
• 29 → 2² + 9² = 85
• 85 → 8² + 5² = 89
• 89 → 8² + 9² = 145
• 145 → 1² + 4² + 5² = 42
• 42 → 4² + 2² = 18
• 18 → 1² + 8² = 65
• 65 → 6² + 5² = 61
• 61 → 6² + 1² = 37
• 37 → 3² + 7² = 58
• 58 → 5² + 8² = 89
• 89 → 8² + 9² = 145

계산 과정을 반복하다 보면, 어느새 같은 수가 나온 뒤로 같은 과정을 반복하고 있는걸 발견할 수 있다. 그에 따라, 계산 과정에서 거쳐가는 수들을 기록해 두었다가, 새 연산을 하기 전에 이미 기록한 수인지 확인하면 행복하지 않은 수를 찾아낼 수 있다.

def happy_number(n):
    """주어진 수가 '행복한 수'인지 확인한다.
    행복한 수는 각 자리 숫자의 제곱의 합을 구하는 연산을 반복해 1이 되는 수이다.

    Args:
        n (int): 행복한 수인지 확인할 자연수

    Returns:
        bool: 행복한 수인지 여부
    """
    trace = set()
    while n not in trace:
        trace.add(r)
        
        sum = 0
        while n > 0:
            sum += (n % 10) ** 2
            n //= 10
        if sum == 1:
            return True
        n = sum
    return False

코드

from sys import stdin

# 에라토스테네스의 체를 사용해 1부터 10,000까지의 소수를 미리 계산한다.
sieve = [True for _ in range(10_001)]
sieve[1] = False

for i in range(2, 100):
    if sieve[i]:
        j = 2
        while i * j <= 10_000:
            sieve[i * j] = False
            j += 1


def happy_number(n):
    """주어진 수가 '행복한 수'인지 확인한다.
    행복한 수는 각 자리 숫자의 제곱의 합을 구하는 연산을 반복해 1이 되는 수이다.

    Args:
        n (int): 행복한 수인지 확인할 자연수

    Returns:
        bool: 행복한 수인지 여부
    """
    trace = set()
    while n not in trace:
        trace.add(r)
        
        sum = 0
        while n > 0:
            sum += (n % 10) ** 2
            n //= 10
        if sum == 1:
            return True
        n = sum
    return False


P = int(stdin.readline().strip())
for _ in range(P):
    idx, M = map(int, stdin.readline().strip().split())
    # M이 소수인지 확인하고. 행복한 수인지 확인한 뒤 둘 다 참이라면 'YES', 그렇지 않다면 'NO'를 출력한다.
    print(f"{idx} {M} {'YES' if sieve[M] and happy_number(M) else 'NO'}")