[Python/BOJ] 4233. 가짜소수

https://www.acmicpc.net/problem/4233

4233번: 가짜소수

문제

페르마의 소정리 (Fermat's little theorem)의 내용은 다음과 같다.

p가 소수일 때, 임의의 정수 a>1에 대해서, ap == a (mod p)가 성립한다.

즉, a를 p제곱한 뒤, p로 나눴을 때, 나머지는 a가 되는 것이다.

하지만, p가 소수가 아닌 경우에 어떤 정수 a에 대해서 위의 식을 만족하는 경우가 있다. 이때, p를 밑이 a인 가짜소수라고 한다. (모든 a에 대해서 식을 만족하는 수를 카마이클 수라고 한다)

p와 a가 주어졌을 때, p가 밑이 a인 가짜소수인지 아닌지 알아내는 프로그램을 작성하시오. 

입력

입력은 여러 개의 테스트 케이스로 이루어져 있다. 각 테스트 케이스는 한 줄로 이루어져 있고, p와 a를 포함하고 있다. 입력의 마지막 줄에는 "0 0"이 주어진다. (2 < p ≤ 1,000,000,000, 1 < a < p)

출력

각 테스트 케이스에 대해서, p가 밑이 a인 가짜소수라면 yes를, 아니라면 no를 한 줄에 하나씩 출력한다.

풀이

Code

import sys
input = sys.stdin.readline

# 1. 소수 판정 리스트 생성 함수 정의
def check_prime_number(num) :
    for i in range(2, int(num**0.5) + 1) :
        if num % i == 0 : return False
    return True

# 2. 제곱 수 생성 함수 정의
def power(a, p, mod) :
    if p == 1 : return a % mod
    x = power(a, p // 2, mod)
    if p % 2 == 0 : return (x * x) % mod
    else : return (x * x * a) % mod
def solution() :
    while True :
        p, a = map(int, input().split())
        # 종료 조건 설정
        if p == a == 0 : break
        # 소수 판별
        if not check_prime_number(p) and power(a, p, p) == a : print('yes')
        else : print('no')

if __name__ == "__main__" :
    solution()