[Python] 17633번 제곱수의 합 (More Huge)

 

 

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

17633번: 제곱수의 합 (More Huge)

---

 

23/12/05

 

 

처음으로 푼 다이아3 문제이다. 높은 수학적 지식을 요구하기 때문에 사실상 혼자 이 문제를 관찰 만을 통해 풀어낸다는 건 불가능에 가깝다고 생각한다.

 

다만 이 문제에서 요구하는 정수론적 지식을 이미 갖추고 있는 사람이라면 다이아 3이라는 난이도에 비해 쉽게 문제를 해결할 수 있을 것이라고 생각한다.

 

---

 

문제 접근 방식:

 

 

먼저, 이 문제는 밀러-라빈 소수 판별법과 폴라드-로 소인수 분해 알고리즘을 구현할 수 있어야 풀 수 있다.

 

밀러-라빈 소수 판별법에 대해 간단하게 설명을 하자면 다음과 같다.

 

밀러-라빈 소수 판별법은 기본적으로 페르마의 소정리에서 출발한다.

 

페르마의 소정리는 아래와 같은 정리다.

$p$가 소수이면, $p$와 서로소인 $a$에 대해서 $a^{p-1} \equiv 1 \textrm{ (mod }p)$이다.

 

이 정리의 대우를 말하면 뒤의 조건을 만족하지 않는 경우에는 합성수라고 말할 수 있다.

 

카마이클 수의 존재 때문에 페르마의 소정리에 대한 역은 항상 참이라고 보장할 수는 없다.

 

따라서 우리는 뒤의 조건을 만족하지 않는 경우, 합성수임을 보장할 수는 있지만 소수임은 확률적으로 보장할 수 있다.

 

따라서 이를 이용하여 확률적으로 소수인지 합성수인지를 가려낼 수 있다.

 

이를 여러 번 반복하면 아주 높은 확률로 소수를 판별할 수 있는데, 특정한 $a$값들에 대해서만 판별하면 PS에서 활용하는 충분히 큰 수 범위 안에서는 확정적으로 소수를 판별할 수 있다.

 

밀러-라빈 소수 판별법의 시간 복잡도는 $\mathcal{O}(k \mathrm{log}^3N$이다.

 

$k$는 판별법을 실행하는 횟수이다.

 

 

폴라드-로 소인수 분해 알고리즘을 간단하게 설명하자면 다음과 같다.

 

특정 다항식(예를 들어 $x^2 + 1$)의 $x$값에 숫자를 넣고 $N$으로 나누는 행위는 $N$이 제한되어 있기 때문에 비둘기집의 원리에 의해서 반복되는 부분이 생길 수밖에 없다.

 

이를 이용하여 $N$의 약수를 찾는데, 플로이드 순환 알고리즘(그냥 두 포인터)을 사용하여 반복되는 부분을 찾아서 약수를 구할 수 있다.

 

이를 소인수가 나올 때까지 반복하면 된다. 소수임을 판단하는 것은 위의 밀러-라빈 소수 판별법으로 빠르게 할 수 있다.

 

 

이 문제에서 사용되는 정리는 3개이다.

 

1. 페르마의 두 제곱수 정리

https://en.wikipedia.org/wiki/Fermat%27s_theorem_on_sums_of_two_squares

Fermat's theorem on sums of two squares - Wikipedia

https://en.wikipedia.org/wiki/Sum_of_two_squares_theorem

Sum of two squares theorem - Wikipedia

어떤 소수 $p$가 두 제곱수의 합으로 표현될 수 있다는 것은 $p$를 $4$로 나눈 나머지가 $1$이라는 것과 동치이다.
즉, 임의의 정수 $n$이 두 제곱수의 합으로 표현될 수 있다는 것은 $n$를 소인수분해 한 결과에서 $p^k$를 포함하고 있지 않다는 것과 동치이다. ($p$는 $4$로 나눈 나머지가 $3$인 소수, $k$는 홀수)

 

2. 르장드르의 세 제곱수 정리

https://en.wikipedia.org/wiki/Legendre%27s_three-square_theorem

Legendre's three-square theorem - Wikipedia

임의의 정수 $n$이 세 개 이하의 제곱수의 합으로 표현될 수 있다는 것은 $n$이 음이 아닌 정수 $a, b$에 대하여 $4^a(8b+7)$꼴로 표현될 수 있다는 것과 동치이다.

 

3. 라그랑주의 네 제곱수 정리

https://en.wikipedia.org/wiki/Lagrange%27s_four-square_theorem

Lagrange's four-square theorem - Wikipedia

임의의 정수 $n$은 4개 이하의 제곱수의 합으로 표현될 수 있다. 

 

 

1번 정리와 2번 정리를 활용하여, 두 개의 제곱수로 나타낼 수 있다면 두 개의 제곱수로 표현한다고 하고, 그렇지 않은 경우 세 개의 제곱수로 표현할 수 있다고 하면 된다. 나머지 경우는 모두 네 개의 제곱수로 표현 가능함이 알려져 있으므로 이를 출력하도록 함수를 구성하면 충분하다.

 

위의 수학적 내용에 대한 증명은 링크를 참고하면 될 것 같다.

  

---

아래는 내가 위의 접근 방식과 같이 작성한 파이썬 코드이다. 더 보기를 누르면 확인할 수 있다.

# 17633번 제곱수의 합 (More Huge)
# 밀러라빈, 폴라드로, 정수론
import sys
input = sys.stdin.readline
from random import randrange
from math import gcd
def Miller_Rabin_primality_test(N, a, s, d):
    if a >= N: a %= N
    if a <= 1: return True
    k = pow(a, d, N)
    if k == 1 or k == N-1: return True
    for r in range(1, s):
        k = pow(k, 2, N)
        if k == N-1: return True
    return False
def is_prime_with_Miller_Rabin(N):
    if N == 2: return True
    if N == 1 or N % 2 == 0: return False
    s, d = 0, N-1
    while d % 2 == 0:
        d //= 2
        s += 1
    if N < 4759123141:
        A = [2, 6, 61]
        for a in A:
            if not Miller_Rabin_primality_test(N, a, s, d):
                return False
        return True
    A = [2, 325, 9375, 28178, 450775, 9780504, 1795265022]
    for a in A:
        if not Miller_Rabin_primality_test(N, a, s, d):
            return False
    return True
def rho(n): # n의 약수를 구한다
    if is_prime_with_Miller_Rabin(n): return n
    if n == 1: return 1
    if n%2 == 0: return 2
    tortoise, c, d = randrange(2, n), randrange(1, n), 1
    hare = tortoise
    while d == 1:
        tortoise = ((tortoise**2 % n) + c)%n
        hare = ((hare**2 % n) + c)%n
        hare = ((hare**2 % n) + c)%n
        d = gcd(abs(tortoise - hare), n)
        if d == n: return rho(n)
    if is_prime_with_Miller_Rabin(n): return d
    else: return rho(d)
def factorization_with_Pollard_rho(N):
    fac = []
    while N > 1:
        d = rho(N)
        fac.append(d)
        N //= d
    fac.sort()
    dic = dict()
    for i in fac:
        if i not in dic:
            dic[i] = 1
        else:
            dic[i] += 1
    return dic
def solve():
    N = int(input())
    factors = factorization_with_Pollard_rho(N)
    ans = 1
    for i in factors.values():
        if i % 2:
            ans += 1
            break
    else:
        return ans
    for p, k in factors.items():
        if k % 2 == 1 and p % 4 == 3:
            ans += 1
            break
    else:
        return ans
    if 2 in factors:
        for i in range(factors[2]//2+1):
            if (N//(4**i))%8 == 7:
                return ans+1
    else:
        if N%8 == 7:
            return ans+1
    return ans
print(solve())