본문 바로가기

컴퓨터 과학/자료구조, 알고리즘

[Rust로 백준 하루 하나] 15-5. 다음 소수

문제 (4134번)

정수 n(0 ≤ n ≤ 4*10^9)가 주어졌을 때, n보다 크거나 같은 소수 중 가장 작은 소수 찾는 프로그램을 작성하시오.

 

입력

첫째 줄에 테스트 케이스의 개수가 주어진다. 각 테스트 케이스는 한 줄로 이루어져 있고, 정수 n이 주어진다.

 

출력

각각의 테스트 케이스에 대해서 n보다 크거나 같은 소수 중 가장 작은 소수를 한 줄에 하나씩 출력한다.


풀이

코드

use std::io::{self, BufReader, BufRead, BufWriter, Write};
use std::f64;
use std::time::Instant;

fn main() {
    let start_time = Instant::now();

    let mut reader = BufReader::new(io::stdin().lock());
    let mut writer = BufWriter::new(io::stdout().lock());

    let mut num = String::new();
    reader.read_line(&mut num).unwrap();
    let num: usize = num.trim().parse().unwrap();

    let mut input = String::new();
    for _ in 0..num {
        reader.read_line(&mut input).unwrap();
        let n: u32 = input.trim().parse().unwrap();

        // n보다 크거나 같은 소수 중 가장 작은 소수
        let mut output = n;
        'a: loop {
            let limit = (output as f64).sqrt() as u32;
            if output < 2 {
                output = 2;
            }
            for i in 2..=limit {
                if output % i == 0 {
                    output += 1;
                    continue 'a;
                }
            }
            break;
        }

        writeln!(writer, "{}", output).unwrap();
        input.clear();
    }

    let end_time = Instant::now();
    let elapsed_time = end_time.duration_since(start_time);
    println!("{:?}", elapsed_time);
}

해설

use std::io::{self, BufReader, BufRead, BufWriter, Write};

fn main() {
    let mut reader = BufReader::new(io::stdin().lock());
    let mut writer = BufWriter::new(io::stdout().lock());

    let mut num = String::new();
    reader.read_line(&mut num).unwrap();
    let num: usize = num.trim().parse().unwrap();

    let mut input = String::new();
    for _ in 0..num {
        reader.read_line(&mut input).unwrap();
        let n: u32 = input.trim().parse().unwrap();

        // n보다 크거나 같은 소수 중 가장 작은 소수
        let mut output = n;
        'a: loop {
            for i in 2..=output/2 {
                if output % i == 0 { break; }
                else if i == output/2 { 
                    break 'a;
                }
            }
            output += 1;
        }

        writeln!(writer, "{}", output).unwrap();
        input.clear();
    }
}

처음엔 이렇게 나누기 2가 되는 수까지 반복을 했으나 시간초과가 나왔습니다. 지난번에 이렇게 풀어서 그대로 했던 건데 사실 제곱근까지만 반복하면 됩니다.

 

제곱근까지 반복하는 방식으로 해결할 수 있었지만 시간을 재보았을 때 3,4초 정도 나왔는데 문제의 시간제한 1초를 어떻게 통과했는지는 잘 모르겠습니다. 아시는 분은 댓글로 알려주시면 감사하겠습니다.

 

Rust에서 as u32를 사용하여 소수점을 포함한 f64 값을 u32로 변환할 때, 소수점 이하 부분은 단순히 버려집니다 (즉, 내림 처리).


추가 학습