백준 “다음 소수” 문제를 풀어보았습니다.

여기를 누르면 건너뛸 수 있습니다.


문제

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

입력

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

출력

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

예제 입력 1

3
6
20
100

예제 출력 1

7
23
101

Solution

[문제 이해]

  • 소수 찾기
  • 입력된 큰 수(40억)와 같거나 보다 큰 소수 찾기 → 에라토스테네스의 체?

이해한 내용을 바탕으로 코드를 작성했습니다.

[코드]

import sys
input = sys.stdin.readline

def solution():
    result = []
    case = int(input())

    # 메모리 사용 효율을 높이기 위해서 제너레이터 사용
    def gen_range(start, end, step=1):
        for i in range(start, end, step):
            yield i

    # 소수인지 확인
    def is_prime(x):
        if x <= 1:
            return False

        m = int(x**0.5)
        for i in range(2, m+1):
            if x % i == 0:
                return False

        return True

    inf = 10**10
    for _ in range(case):
        n = int(input())
        prime = None
        # 입력된 n ~ inf 중에 소수 찾기
        for i in gen_range(n, inf):
            if is_prime(i):
                prime = str(i)
                break
        result.append(prime)

    print('\n'.join(result))

solution()

[과정]

에라토스테네스의 체로 작성한 코드는 시간 초과로 실패했습니다.

그래서 단순하게 입력된 1개의 수가 소수인지 확인하며 무한으로 수를 키워가는 방식으로 접근했습니다.