BOJ 17633 - 제곱수의 합 (More Huge) (Python3)

모든 시험이 오늘부로 종료되었다. 성적이 나온 과목은 몇 개 있는데 하나는 생각보다 덜 만족스러웠다. 어쩌겠는가, 이미 안다고 자만하며 공부를 조금 소홀히 한 내 잘못이지.

이번 문제는 정수론적 지식이 매우 강력하게 필요한, 무려 다이아몬드 IV 문제이다. 소인수분해 과정에서 밀러-라빈과 폴라드-로까지 사용해야 하는 매우 복잡한 문제. 혹시 그 쪽에 대한 지식이 필요하다면, 이 포스트를 참조하라.

심지어 이 문제를 약간 발전시킨 문제는 루비 IV에 랭크되어 있다! 그 코드도 거의 다 만들었으나, 의문의 시간초과로 인해 난항이다.

 

그럼 본격적으로 풀이에 돌입해 보자.

---

문제

라그랑주는 1770년에 모든 자연수는 넷 혹은 그 이하의 제곱수의 합으로 표현할 수 있다고 증명하였다. 어떤 자연수는 복수의 방법으로 표현된다. 예를 들면, 26은 5²과 1²의 합이다; 또한 4²+3²+1²으로 표현할 수도 있다.

역사적으로 암산의 명수들에게 공통적으로 주어지는 문제가 바로 자연수를 넷 혹은 그 이하의 제곱수 합으로 나타내라는 것이었다.

1900년대 초반에 한 암산가가 15663=125²+6²+1² +1²라는 해를 구하는데 8초가 걸렸다는 보고가 있다. 좀 더 어려운 문제에 대해서는 56초가 걸렸다: 11339=105²+15²+8²+5².

자연수 n이 주어질 때, n을 최소 개수의 제곱수 합으로 표현하는 컴퓨터 프로그램을 작성하시오.

 

입력

입력은 표준입력을 사용한다. 입력은 자연수 n을 포함하는 한 줄로 구성된다. 여기서, 1 ≤ n ≤ 1,000,000,000,000,000,000이다.

 

예제 입력)

# 예제 입력 1
25

# 예제 입력 2
26

# 예제 입력 3
11339

# 예제 입력 4
34567

 

출력

출력은 표준출력을 사용한다. 합이 n과 같게 되는 제곱수들의 최소 개수를 한 줄에 출력한다.

 

예제 출력)

# 예제 출력 1
1

# 예제 출력 2
2

# 예제 출력 3
3

# 예제 출력 4
4

 

---

내 코드

from random import randrange
pli = [2,3,5,7,11,13,17,19,23,29,31,37,41,43]

# Miller-Rabin Test
def power(x,y,p):
    if y<2: return (x**y)%p
    else:
        d = y//2
        if y%2 == 0: return (power(x,d,p)**2)%p 
        else: return (x*(power(x,d,p))**2)%p
def miller_rabin(n,a):
    r,d = 0,n-1
    while d%2 == 0: r += 1; d //= 2
    x = power(a,d,n)
    if x == 1 or x+1 == n: return True
    for i in range(0, r-1):
        x = power(x,2,n)
        if x+1 == n: return True
    return False
def isprime(n):
    if n in pli: return True
    if n == 1 or n%2 == 0: return False
    for p in pli:
        if not miller_rabin(n,p): return False
    return True

# Pollad-Rho algorithm
def gcd(a,b):
    if a<b: a,b = b,a
    while b != 0: a,b = b,a%b
    return a
def rho(n):
    if isprime(n): return n
    if n == 1: return 1
    if n%2 == 0: return 2
    x,c,d = randrange(2,n),randrange(1,n),1
    y = x
    while d == 1:
        x = ((x**2%n)+c)%n
        y = ((y**2%n)+c)%n
        y = ((y**2%n)+c)%n
        d = gcd(n,abs(x-y))
        if d == n: return rho(n)
    if isprime(d): return d
    else: return rho(d)

# Answer Code
def is4(n):
    while True:
        if n%4 == 0: n //= 4
        else: return True if n%8 == 7 else False
def issquare(n):
    return True if n==int(n**.5)**2 else False
n = int(input())
if issquare(n): print(1)
elif is4(n): print(4)
else:
    data = {}
    while n != 1:
        d = rho(n)
        try: data[d] += 1
        except: data[d] = 1
        n //= d
    for k,v in data.items():
        if k%4 == 3 and v%2 == 1: print(3); quit()
    print(2)

 

N이 충분히 작으면 다이나믹 프로그래밍이나 하나씩 다 대입해 보는 브루트 포스로도 풀 수 있다. 그런 문제도 이미 있으며, 실버 정도의 난이도면 충분.

그러나 이 경우에는 N이 \(10^{18}\)으로, 비상식적인 크기를 자랑한다. 당연히 하나하나 대입하는 무식한 방법으로 풀다간 며칠간 컴퓨터를 돌려야 하는 불상사에 직면하게 된다.

정수론적 지식을 크게 요구한다는 것이 바로 이 지점. 상당히 심화된 정수론 지식을 이용해서 풀어야 한다.

여기에 있는 모든 정리의 증명은 따로 포스트를 발행하여 그 곳에서 담겠다. 증명의 난이도가 제법 있으며, 지면은 너무 적다!

 

문제에서 힌트를 줬듯이, 라그랑주의 네 제곱수 정리에 의해 답은 1, 2, 3, 4 중 하나이다. 이 중 답이 1인 경우야 해당하는 수가 제곱수이니 issquare 함수로 편하게 거르면 된다.

그 다음으로 답이 4인 경우를 거를 수 있는데, 이 경우는 정수론적 지식이 크게 필요하다. 르장드르의 세 제곱수 정리에 의하면, \(4^m(8k+7)\)꼴이 아닌 모든 정수는 제곱수 세 개의 합으로 나타낼 수 있다고 알려져 있다. 그 경우를 is4 함수로 거르면 끝.

그리고 걸러지지 않은 모든 정수들은 답이 2 아니면 3일 것이다. 그렇다면 어떻게 이 둘을 분리할 것인가?

답은 소인수분해다. 여기서 밀러-라빈과 폴라드-로가 튀어나오는 것이다!

 

페르마의 두 제곱수 정리에 의해서, 소인수분해를 했을 때 \(4k+3\) 꼴의 소인수의 지수가 모두 짝수여야만 제곱수 2개의 합으로 나타낼 수 있음이 알려져 있다.

그렇다면 폴라드-로 알고리즘을 통해 소인수분해를 한 뒤, 그 결과를 data에 {소인수: 지수} 꼴로 저장하자. \(2^4\cdot3^7\cdot7^2\)를 {2:4, 3:7, 7:2} 꼴의 딕셔너리로 저장을 하자는 것이다.

그 다음 모든 data의 요소를 순회하며 소인수가 \(4k+3\) 꼴일 때 지수가 짝수인지 홀수인지 판별하고, 그 과정에서 홀수가 한 번이라도 나오면 3을, 그렇지 않으면 2를 출력하도록 코드를 짜면 끝이다!

 

이로서 17633번의 풀이를 마친다.

그럼, 오늘도 당신의 코딩 실력이 상승하기를.