[백준/Python] 17626 : Four Squares

 

시간 제한이 굉장히 타이트한 문제이다.

처음엔 매번 num에서 가능한 제곱수를 찾아 빼는 방식으로 작동시켰었는데, 이는 큰 수에 대해 매우 비율적이기 때문에 dp 배열을 만들어 각 숫자를 만들 수 있는 제곱수의 최소 개수를 저장하도록 작성하였다.

 

동적 계획법(DP) ?

- 각 숫자를 만들 수 있는 최소 제곱수의 개수를 저장하고, 그 값을 이전에 계산된 결과를 활용해 구하는 방법

 

import math

num = int(input())
dp = [float('inf')]*(num + 1)
dp[0] = 0 # 0을 만들기 위한 최소 제곱수 개수는 0

for i in range(1, num + 1):
    for j in range(1, int(math.sqrt(i))+ 1):
        dp[i] = min(dp[i], dp[i-j*j] + 1)
print(dp[num])

 

- 먼저 dp[i]를 숫자 i를 만들 수 있는 최소 제곱수의 개수로 정의한다.

- 0은 제곱수가 필요 없기 때문에 초기 값으로는 dp[0] = 0을 설정한다.

- 나머지 dp[i] 는 dp[i - j^2] + 1로 업데이트한다. (j는 가능한 제곱수)

 

바깥 반복문 : 1부터 num 까지 각 숫자에 대해 최소 제곱수 개수를 계산

안쪽 반복문 : i보다 작거나 같은 제곱수 j^2를 찾아 dp[i - j^2] 값을 갱신

 

이렇게 작성하면 시간 복잡도는  O(n*sqrt(n))으로, 각 숫자 i에 대해 제곱수를 확인하는데 최대 sqrt(n) 번의 연산이 필요하게 되어 짧은 시간 안에도 통과할 수 있게 된다.

 

사실 처음에 시간초과 오류가 많이 났었고, 심지어 위 코드로도 시간초과가 나서 갈피를 못잡았었는데, 찾아보니 python3가 아닌 pypy3으로 바꿔서 제출하면 통과될수도 있다고 하기에 바꿔봤더니 바로 통과되었다.