[BOJ, 백준] 4320 완전 P제곱수 python

오늘은 밀러 - 라빈 , 폴라드 - 로의 하루가 될 것 같은 느낌.

골드 5였던 문제였고, 주의해야할 점이 상당히 많았던 문제.

 

아래는 밀러-라빈 / 폴라드-로때 작성한 코드

import sys
import random
from collections import defaultdict
import math

sys.setrecursionlimit(10 ** 6)
input = sys.stdin.readline

def power(x, y, p):
    res = 1
    x = x % p
    while y > 0:
        if y & 1:
            res = (res * x) % p
        y = y >> 1
        x = (x * x) % p
    return res

def gcd(a, b):
    if a < b:
        a, b = b, a
    while b != 0:
        r = a % b
        a = b
        b = r
    return a

def miller_rabin(n, a):
    r = 0
    d = n - 1
    while d % 2 == 0:
        r += 1
        d = d // 2

    x = power(a, d, n)
    if x == 1 or x == n - 1:
        return True

    for i in range(r - 1):
        x = power(x, 2, n)
        if x == n - 1:
            return True
    return False

def is_prime(n):
    alist = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41]
    if n == 1:
        return False
    if n == 2 or n == 3:
        return True
    if n % 2 == 0:
        return False
    for a in alist:
        if n == a:
            return True
        if not miller_rabin(n, a):
            return False
    return True


def pollardRho(n):
    if is_prime(n):
        return n

    if n == 1:
        return 1

    if n % 2 == 0:
        return 2

    x = random.randrange(2, n)
    y = x
    c = random.randrange(1, n)
    d = 1
    while d == 1:
        x = ((x ** 2 % n) + c + n) % n
        y = ((y ** 2 % n) + c + n) % n
        y = ((y ** 2 % n) + c + n) % n
        d = gcd(abs(x - y), n)
        if d == n:
            return pollardRho(n)
    if is_prime(d):
        return d
    else:
        return pollardRho(d)

 

 

 

아래 부분이 해당 알고리즘을 활용한 본 문제의 실제 코드가 되겠다.

 

while True:
    x = int(input())
    if x == 0:
        break

    prime = defaultdict(int)
    y = abs(x)
    while y > 1:
        divisor = pollardRho(y)
        prime[divisor] += 1
        y = y // divisor

    degree = list(set(prime.values()))
    degree_gcd = math.gcd(*degree)

    for deg in degree:
        deg //= degree_gcd

    if x < 0:
        if degree_gcd % 2 == 0:
            while degree_gcd % 2 == 0:
                degree_gcd //= 2
    
    print(degree_gcd)