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

내가 이 블로그를 개설할 때만 해도, Ruby 카테고리는 구색맞추기용이었다. 내가 내 실력으로 Ruby 난이도의 문제를 하나라도 풀 수 있을 것이라는 생각을 한 적은 없었다.

작년 말에 상황이 달라졌다. 적당한 알고리즘을 검색했고, 이렇게 하면 풀 수 있을지도 모르겠다는 생각이 들었다. 길이 희미하게 보이기 시작했다.

1월에 제대로 코드를 짜고 도전해 보기 시작했다. 계속되는 시간초과에 발이 묶여 여기저기 도움을 요청해 봤지만, 이렇다 할 답은 나오지 않았다. 현 시점에서 푼 사람이 69명뿐인 문제이니 그럴 만도 하지.

그리고 오늘. 6월 6일 새벽 1시쯤에 번쩍 하고 스쳐지나가는 아이디어가 있었고, 그걸 구현하려다가 빼먹은 게 좀 있어서 몇 번 틀려 준 다음 올라가는 퍼센테이지를 보고 정말 울 뻔 했다. 이게 진짜 될 줄 몰랐어서. 그리고 너무 허무하게 무너져서.

다섯 달 동안 찾아헤메던 열쇠가 주머니 속에 있음을 깨닫는 데는 채 1시간도 걸리지 않았으니.

 

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

---

문제

라그랑주는 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과 같게 되는 제곱수들의 최소 개수를 첫째 줄에 출력한다.

둘째 줄에는, 제곱의 합이 n과 같게 되는 수들을 첫째 줄에 출력한 갯수만큼 공백으로 구분하여 출력한다. 음의 정수는 출력해서는 안된다.

답이 여러개인 경우, 아무거나 출력해도 좋다.

 

예제 출력)

# 예제 출력 1
1
5

# 예제 출력 2
2
1 5

# 예제 출력 3
3
1 27 103

# 예제 출력 4
4
1 22 109 149

 

---

내 코드

# Deleted #

 

UPD) 코드를 더 이상 공개하지 않는다. 이 글을 쓴 이후, 내 코드를 주석까지 복사 후 붙여넣기한 사람을 세 명이나 목격했다. 솔직히 이 정도 난이도가 되는 문제는 코드를 공개하는 의미가 없다고 본다. 풀이 방법만 확인 후 자기 손으로 구현할 줄은 알아야 한다고 생각한다. 무슨 자격이 있어야 이런 주제넘는 말을 할 수 있을지는 모르겠지만, 일단 치팅의 피해자라면 발언권은 있다고 본다.

 

진짜 끔찍하게 길다... gcd, miller_rabin, isprime, rho, power, is4, issquare, is3의 함수는 17633번 문제나 다른 문제 풀이과정에서 이미 사용했고 작동 방식도 설명했으니 넘어가도록 하겠다. 이 글에서 중요한 것은 궁극적으로는 for2, for3 함수, 그리고 그것들을 도출하는 데 도움을 주는 tonelli, cornacchia가 있겠다.

답을 얻기 위해서 가장 중요한 것은 무엇인가? 어떤 제곱수의 합으로 이 수가 나타내어지는지 정하는 것이 먼저겠지. 당연히 제곱수는 답이 너무 자명하니 가장 먼저 걸러준다

네 개의 제곱수 합으로 표현할 수 있는 n은 무조건 \(n=4^N(8k+7)\) 꼴이다.

이 사실을 이용하면 \(4^N(8k+6)=a^2+b^2+c^2\)일 시 \(n=a^2+b^2+c^2+(2^N)^2\)이게 되므로, 어떤 수가 두 제곱수 혹은 세 제곱수의 합으로 표시되는 형태를 구할 수 있다면, 네 제곱수 케이스 역시도 가능함을 알 수 있다.

그렇다면 세 제곱수 케이스는 어떻게 풀어야 할까?

 

Landau와 Ramanujan의 연구 결과에 의하면, \(x\) 이하의 자연수 중 두 제곱수의 합으로 나타내어지는 자연수의 수를 \(N(x)\)라고 두면 어떠한 상수 \(b\)에 대해서 \(\lim_{x\to\infty}\frac{N(x)\sqrt{\log x}}{x}=b\)이 성립한다고 한다(이 수는 그 두 명의 이름을 따서 Landau-Ramanujan constant라는 이름이 붙어 있다.). 당연히 \(\sqrt{\log x}\)는 \(x\)에 비견해서는 매우 작으니, \(N(x)\)는 \(x\)와 크게 다르지 않다는 주장은 말이 된다.

그렇다면 대부분의 자연수는 두 제곱수의 합으로 표시가 된다는 말이다. 그말인즉슨, \(N\)이 세 제곱수의 합으로 나타내어질 경우, 여러 번 시도해 보면 두 제곱수의 합으로 나타내어지는 \(N-a^2\)을 그렇게 느리지 않은 시간 내에 찾을 수 있다는 소리가 된다. 나는 \(a\)를 \(\lfloor\sqrt{N}\rfloor\)부터 시작하여 하나씩 줄여가는 방식으로 계산했다.

이 말 그대로 구현하면 시간초과가 날 것이다. 만약 \(n=4^mr\)이고 \(r\)이 4의 배수가 아니라면, \(r=a^2+b^2+c^2\)인 세 정수를 구한 다음 \(n=(2^ma)^2+(2^mb)^2+(2^mc)^2\)로 세 개의 제곱수를 구해도 큰 문제가 없다. 이렇게 하지 않으면 어째서인지는 몰라도 \(3\cdot 2^{50}\)과 같은 수에서 시간초과가 나더라.

자. 네 제곱수 케이스를 세 제곱수 케이스로, 세 제곱수 케이스를 두 제곱수 케이스로 환원하는 방법을 설명하였다. 이제 남은 것은 무엇? 두 제곱수 케이스를 공략하는 것이다.

 

이 때 등장하는 알고리즘이 바로 Cornacchia's algorithm이다. 이 알고리즘은 \((d,m)=1, 0<d<m\)일 때 \(x^2+dy^2=m\)의 정수해를 구하는 알고리즘인데, 여기서 \(d=1,m=p\)로 두면 목표로 했던 \(x^2+y^2=p\)의 정수해를 구할 수 있게 된다!

그 알고리즘의 풀이 방법 중 첫 단계가 \(r_0^2\equiv -1 \pmod p\)인 정수 \(r_0\)을 찾는 것인데, 물론 그러한 \(r_0\)은 존재하겠지만 구하는 것이 쉽지는 않다. \(p\)가 매우 작으면 윌슨 정리를 응용해서 직접 계산할 수 있지만 또 여기서의 \(p\)는 18자리까지, 무식하게 커질 수 있다는 것을 고려해 주어야 한다.

이 이산제곱근을 구하려면 또 다른 알고리즘을 가지고 와야 하는데, 그 알고리즘이 바로 Tonelli-Shanks algorithm이다. 다른 알고리즘은 모두 위키백과를 보고 직접 구현했는데 이건 도저히 못 알아먹겠어서 남의 구현체를 보고 이해하는 기간을 거쳤다.

여하튼! 이렇게 구한 소수들의 두 제곱수 표현들을 \((a^2+b^2)(c^2+d^2)=(ac+bd)^2+(ad-bc)^2\)를 이용해서 원래 \(n\)에 대한 두 제곱수 표현을 만들어낼 수 있다!

0.5초에 추가시간도 없는 문제라 Python3으로 공략하기 많이 힘들 것 같다고 생각할 수 있으나, 내 코드가 대략 176ms에 돌아간다. 시간에 대한 걱정만큼은 안 해도 된다.

 

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

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