[백준/Python3] 17626번 Four Squares

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

17626번: Four Squares

 

---

문제

라그랑주는 1770년에 모든 자연수는 넷 혹은 그 이하의 제곱수의 합으로 표현할 수 있다고 증명하였다. 어떤 자연수는 복수의 방법으로 표현된다. 예를 들면, 26은 5^2과 1^2의 합이다; 또한 4^2 + 3^2 + 1^2으로 표현할 수도 있다. 역사적으로 암산의 명수들에게 공통적으로 주어지는 문제가 바로 자연수를 넷 혹은 그 이하의 제곱수 합으로 나타내라는 것이었다. 1900년대 초반에 한 암산가가 15663 = 125^2 + 6^2 + 1^2 + 1^2라는 해를 구하는데 8초가 걸렸다는 보고가 있다. 좀 더 어려운 문제에 대해서는 56초가 걸렸다: 11339 = 105^2 + 15^2 + 8^2 + 5^2.

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

---

입력

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

---

출력

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

---

문제풀이

먼저 처음에는 for문으로 n의 제곱근부터 1까지, 큰 수부터 제곱하여 n에서 순차적으로 빼면 가능할 것이라고 생각해 아래와 같은 코드를 작성하였다. 

n = int(input())
ns = int(n**(1/2)) #n의 제곱근 중 가장 큰 정수
ans = 0

for i in range(ns, 0, -1):
    if n == 0:
        break
    if n < 4:
        ans += n
        break
    while i**2 <= n:
        ans += 1
        n -= i**2

print(ans)

하지만 32와 같은 반례를 발견하여 아래와 같이 수정하였다. while문을 두고 만약 ans값이 4보다 크다면, 가장 큰 수가 아닌 그 이후 수가 네 제곱수 중 가장 큰 수일 것이기 때문에 for문의 시작인 ns를 1을 줄여주고 다시 처음부터 시작한다.

이렇게 하면 32의 경우 위 코드에서 32=5^2+2^2+1+1+1로 ans가 5가 나오지만, 아래 코드에선 32=4^2+4^2로 2가 정상적으로 나와 문제의 조건인 최소 개수의 제곱수의 합으로 표현된다. 

n = int(input())
n_save = n
ns = int(n**(1/2))
ans = 0

while 1:
    for i in range(ns, 0, -1):
        if n == 0:
            break
        if n < 4:
            ans += n
            break
        while i**2 <= n:
            ans += 1
            n -= i**2

    if ans <= 4:
        break
    else:
        n = n_save
        ans = 0
        ns -= 1


print(ans)

하지만 이렇게 해도 틀려서 새로운 반례를 찾던 중 72와 같은 반례를 찾게 되었다. 72=8^2+2^2+2^2로 3이 나오지만, 72=6^2+6^2로 2라는 최소 개수의 제곱수의 합이 존재하기 때문에 틀렸다는 것을 알게되었다.

while문으로 큰 수부터 내려오는 경우로 조건을 설정하기에는 점차 복잡함을 느껴 새로운 방식을 찾게 되었는데, 그것은 1부터 시작하여 각 수마다의 최소 개수의 제곱수의 합을 구해놓고 점차 큰 수의 최소 개수의 제곱수의 합을 구하는 방식을 생각하였다.

n = int(input())
ans = [0, 1] #0과 1일 때의 값을 저장한다.

for i in range(2, n+1):
    m = 5 #ans의 원소는 1~4의 값을 띌 것이 분명하기 때문이다.
    
    for j in range(1, int(i**(1/2))+1):
        m = min(m, ans[i-j**2])

    ans.append(m+1)

print(ans[-1])

첫 for문의 i에서 그보다 작은 제곱수들(j)을 뺏을 때의 값이 최소인 수와 그 j^2의 수가 합쳐졌을 때 i가 나올 것이다. 즉 ans[i-j**2]의 값에서 빼준 j라는 수를 더하여 m+1을 n번째에 append하면서 ans를 구하고 그 배열 ans의 끝 값을 출력하면 문제가 원하는 답을 도출할 수 있다.

하지만 이 문제는 계속 시간 초과 문제가 발생하여 pypy3으로 제출하였다. 질문 게시판을 확인해본 결과, 답변으로도 pypy3로 제출하라고 하였다. 구글링을 했을 땐 브루트포스로 풀었을 시 문제가 없다고 한다. 하지만 DP로 풀어버린 내 사고방식을 작성하기 위해 여기서 글을 마친다.