[Python] 백준 - 1699 제곱수의 합

출처: 백준 온라인 저지
https://www.acmicpc.net/problem/1699

처음 풀이 (시간 초과)

from sys import stdin


n = int(stdin.readline())

sq_sum = [n] * (n + 1)
sq_sum[1] = 1

for i in range(2, n + 1):
    if int(i ** 0.5) == i ** 0.5:
        sq_sum[i] = 1
    else:
        for j in range(1, i // 2 + 1):
            if sq_sum[i] > sq_sum[i - j] + sq_sum[j]:
                sq_sum[i] = sq_sum[i - j] + sq_sum[j]
            if sq_sum[i] == 2:
                break

print(sq_sum[n])

자연수 n을 제곱수들의 합으로 표현할 때 그 항의 최소개수를 sq_sum(n)이라고 하자. n이 제곱수인 경우, sq_sum(n)은 1이다. n이 제곱수가 아닌 경우, sq_sum(n)은 sq_sum(n - k) + sq_sum(k) 중 가장 작은 값이다.

i) n이 제곱수인 경우 -> sq_sum(n) = 1
ii) n이 제곱수가 아닌 경우 -> sq_sum(n) = min(sq_sum(n - k) + sq_sum(k))

처음 풀이에서는 위의 점화식을 계산하기 위해, k에 가능한 값을 모두 대입해보는 완전탐색 방식으로 코드를 구현해보았다. 시간 복잡도를 조금이라도 줄이기 위해, sq_sum[n]이 2인 경우 반복문을 나가도록 하였다. (n이 제곱수가 아닌 경우, sq_sum[n]이 가질 수 있는 가장 작은 값은 2이기 때문이다.) 하지만 시간 복잡도가 높은 완전탐색으로 구현한 코드는 예상대로 시간 초과가 떴다.

통과한 풀이

from sys import stdin


n = int(stdin.readline())

sq_sum = [n] * (n + 1)
sq_sum[1] = 1

square = []

for i in range(1, int(n ** 0.5) + 1):
    square.append(i * i)

for i in range(2, n + 1):
    if i in square:
        sq_sum[i] = 1
    else:
        for j in square:
            if j >= i:
                break
            if sq_sum[i] > sq_sum[i - j] + sq_sum[j]:
                sq_sum[i] = sq_sum[i - j] + sq_sum[j]
            if sq_sum[i] == 2:
                break

print(sq_sum[n])

처음 코드에서는 sq_sum[i] = sq_sum[i - j] + sq_sum[j]의 점화식에서 j에 가능한 모든 값(1 ~ i//2 + 1)을 대입했다. 그런데 생각해보니, 어차피 sq_sum[n]에서 n을 제곱수들의 합이므로, j에 제곱수만을 대입해도 모든 경우의 수를 다 포괄할 수 있었다. 그래서 n 이하의 제곱수들의 목록을 square 리스트에 저장한 뒤, j에는 i보다 작은 square 리스트의 값만을 대입하도록 코드를 수정했다. 시간 초과 문제가 해결되었다.

최종 풀이

from sys import stdin


n = int(stdin.readline())

sq_sum = [n] * (n + 1)
sq_sum[1] = 1

square = []

for i in range(int(n ** 0.5) + 1):
    square.append(i * i)

for i in range(2, n + 1):
    if i in square:
        sq_sum[i] = 1
    else:
        # a^2 <= i < (a+1)^2 (a는 i 이하의 가장 큰 제곱수)
        a = int(i ** 0.5)

        # j는 1부터 a까지만 순회함
        for j in square[1:a + 1]:
            if sq_sum[i] > sq_sum[i - j] + sq_sum[j]:
                sq_sum[i] = sq_sum[i - j] + sq_sum[j]
            if sq_sum[i] == 2:
                break

print(sq_sum[n])

sq_sum[i] = sq_sum[i - j] + sq_sum[j] 식에서 j는 제곱수 중에서도 i보다 작은 값만 대입할 수 있다. 그래서 2번째 풀이에서는 for j in square 반복문마다 j가 i보다 작은지를 검사했는데, 이는 비효율적이다. 따라서 최종 풀이에서는 i 이하의 가장 큰 제곱수를 미리 변수 a에 대입한 뒤, 반복문을 square의 a 인덱스까지만 돌게 하여 효율성을 높였다.