[백준 / Rust] 10434번: 행복한 소수

문제

10434번: 행복한 소수

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

313 331 367 ...

경복 : ??

강산 : 379.

경복 : 뭐?

강산 : 행복한 소수잖아.

경복 : 행복한 뭐?

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

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

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

7은 행복한 수이다 ☺.

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

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

입력

첫 줄에 테스트 케이스의 수 P가 주어진다. (1 ≤ P ≤ 1000)

각 테스트 케이스는 테스트 케이스 번호와 행복한 소수인지 판정해야 할 정수인 M으로 이루어져 있다. (1 ≤ m ≤ 10000).

출력

각 테스트 케이스마다, 테스트 케이스의 번호, 입력받은 수, 만일 M이 행복한 소수라면 YES 아니라면 NO를 공백으로 각각 구분하여 출력한다.

---

문제 이해하기

1. 1~10000 사이의 수들 중 어떤 수가 소수인지 에라토스테네스의 체를 사용해서 판별하기
2. 1~10000 사이의 수들 중 소수이면서 행복한 수 찾기
3. 입력으로 들어온 m이 행복한 소수인지 출력

 

 여기서 2번 행복한 수 찾기는 각 자릿수의 제곱 합이 1이 아니면서 사이클이 발견되는 경우에 행복한 수가 아니라고 볼 수 있습니다. 예를 들어, 2는 2, 4, 16, 37, 58, 89, 145, 42, 20, 4... 순서로 진행되는데 여기서 4가 중복으로 발견되었으니 여기서부터 또 다른 사이클이 시작된다고 볼 수 있습니다. 

---

코드

use std::io::{stdin, stdout, BufWriter, Write};

fn main() {
    let mut writer = BufWriter::new(stdout());
    let mut lines = stdin().lines();

    let mut is_prime = vec![true; 10001];
    is_prime[0] = false;
    is_prime[1] = false;
    let mut i = 2;
    while i * i <= 10000 {
        if is_prime[i] {
            let mut j = i * i;
            while j <= 10000 {
                is_prime[j] = false;
                j += i;
            }
        }
        i += 1;
    }

    let mut is_happy_prime = vec![true; 10001];
    is_happy_prime[0] = false;
    is_happy_prime[1] = false;

    for i in 2..=10000 {
        if is_prime[i] && is_happy_prime[i] {
            let mut n = i;
            let mut sum = 0;
            while n > 0 {
                sum += (n % 10) * (n % 10);
                n /= 10;
            }
            if sum == 1 {
                continue;
            }
            let mut visited = vec![false; 10001];
            visited[i] = true;
            while sum != 1 {
                if visited[sum as usize] {
                    is_happy_prime[i] = false;
                    break;
                }
                visited[sum as usize] = true;
                n = sum;
                sum = 0;
                while n > 0 {
                    sum += (n % 10) * (n % 10);
                    n /= 10;
                }
            }
        } else {
            is_happy_prime[i] = false;
        }
    }

    let p = line_to_number(&lines.next().unwrap().unwrap());

    for _ in 0..p {
        let (n, m) = {
            let nums = split_line_to_numbers(&lines.next().unwrap().unwrap());
            (nums[0], nums[1])
        };

        writeln!(
            writer,
            "{} {} {}",
            n,
            m,
            if is_happy_prime[m as usize] {
                "YES"
            } else {
                "NO"
            }
        )
        .unwrap();
    }
}

fn split_line_to_numbers(s: &str) -> Vec<i32> {
    s.split_whitespace().map(|s| s.parse().unwrap()).collect()
}

fn line_to_number(s: &str) -> i32 {
    s.parse().unwrap()
}

시간: 8 ms

메모리: 13216 KB
분류: 수학, 구현, 정수론, 시뮬레이션, 소수 판정

글에서 수정이 필요한 부분이나 설명이 부족한 부분이 있다면 댓글로 남겨주세요!