[Python] 백준/BOJ 17626번: Four Squares (Silver 3)

💻 Problem

문제 보러 가기

 

라그랑주는 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을 최소 개수의 제곱수 합으로 표현하는 컴퓨터 프로그램을 작성하시오.

 

💡 Approach

더 작은 제곱수의 합으로 더 큰 제곱수의 합을 만들 수 있으니 dp로 풀어야겠다는 생각은 들었다.

근데 점화식을 어떻게 만들어야 할지 모르겠어서 일단 다 적어보았다.. (틀린 방법이니 참고 ㄴㄴ)

• dp[1] = 1 (1)
• dp[2] = 2 (1 + 1)
• dp[3] = 3 (1 + 1 + 1)
• dp[4] = 1 (2^2)
• dp[5] = 2 = dp[4] + 1 (2^2 + 1)
• dp[6] = 3 = dp[5] + 1 (2^2 + 1 + 1)
• dp[7] = 4 = dp[6] + 1 (2^2 + 1 + 1 + 1)
• dp[8] = 2 = dp[4] + dp[4] (2^2 + 2^2)
• dp[9] = 1 (3^2)
• dp[10] = 2 = dp[9] + 1 (3^2 + 1)
• dp[11] = 3 = dp[10] + 1 (3^2 + 1 + 1)
• dp[12] = 3 = (2^2 + 2^2 + 2^2)
• dp[13] = 2 = dp[9] + dp[4]
• ...

위처럼 추측하며 점화식을 세워보니 아래 점화식이 나왔다.

dp[i] = min(dp[i], dp[j] + dp[i - j])

 

먼저 index를 value로 초기화한 dp 배열을 만들고 index가 제곱수인 칸은 1로 초기화했다.

세워둔 점화식으로 이중 반복문을 돌렸다.

import sys
input = sys.stdin.readline

n = int(input())
dp = [i for i in range(n + 1)]

for i in range(2, int(n ** 0.5) + 1):
    dp[i ** 2] = 1

for i in range(2, n + 1):
    for j in range(1, i // 2 + 1):
        dp[i] = min(dp[i], dp[j] + dp[i - j])

print(dp[-1])

근데 이러면 시간복잡도가 O(n^2)이 된다.

문제의 시간제한은 0.5초인데......

n의 최댓값은 5 * 10 ^ 4이다.

 

어떻게 최적화할 수 있을지 찾아봤다.

i를 만들 수 있는 제곱수를 기준으로 하여 여기에 i가 될 수 있도록 dp를 활용해 더한다.

• dp[1] = 1 (1)
• dp[2] = 2 (1 + 1)
• dp[3] = 3 (1 + 1 + 1)
• dp[4] = 1 (2^2)
• dp[5] = 2 = dp[5 - 4] + 1 (1+ 2^2)
  • dp[1] + 1에서 + 1이 i(5)를 만들 수 있는 기준 제곱수(2^2)를 나타낸다.
• dp[6] = 3 = dp[6 - 4] + 1 (1 + 1 + 2^2)
• dp[7] = 4 = dp[7 - 4] + 1 (1 + 1 + 1 + 2^2)
• dp[8] = 2 = dp[8 - 4] + 1 (2^2 + 2^2)
• dp[9] = 1 (3^2)
• dp[10] = 2 = dp[10 - 9] + 1 (1 + 3^2)
• dp[11] = 3 = dp[11 - 9] + 1 (1 + 1 + 3^2)
• dp[12] = 3 = dp[12 - 4] + 1 = (2^2 + 2^2 + 2^2)
• dp[13] = 2 = dp[13 - 9] + 1 = (2^2 + 3^2)
• ...

내가 처음에 접근한 방식이랑 어떻게 보면 좀 반대로(?) 접근해야 했다.

dp 어려웡 ㅠ

 

아무튼 다시 정리한 전개를 보면 점화식은 min(dp[i], dp[i - j * j] + 1)이 된다.

전체 로직은 다음과 같다.

1. i를 만들 수 있는 제곱수를 기준으로 한다.
   1. 1부터 시작하여 하나씩 증가하는 j가 있을 때, j * j가 그 제곱수이다.
2. i에서 제곱수를 뺀 dp 값과 1(제곱수 하나)을 더하면 i를 만들 수 있는 제곱수의 합 구성 개수이다.
   1. j 값을 늘리며 제곱수의 합 구성 개수의 최솟값을 갱신한다.

 

근데 모든 자연수가 넷 이하의 제곱수 합으로 표현할 수 있다는 게 신기하네😲

 

✏️ Solution

import sys
input = sys.stdin.readline

n = int(input())
dp = [i for i in range(n + 1)]

for i in range(4, n + 1):
    j = 1
    while j * j <= i:
        dp[i] = min(dp[i], dp[i - j * j] + 1)
        j += 1

print(dp[-1])

바깥 반복문을 n번만큼 반복하고, 안쪽 반복문을  √n번만큼 반복한다.

• 시간복잡도: O(n√n)

n의 최댓값은 5 * 10^4이다.

대입하면 약 5√5 * 10^6 => 약 10^7이다.

시간제한 0.5초는 5 * 10^7이므로 아슬아슬하게 통과한다.

 

python으로 제출하면 시간초과이다.

pypy로 제출해야 통과한다.