[Python] 1699번 제곱수의 합 / 17626번 Four Squares

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

1699번: 제곱수의 합

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

17626번: Four Squares

---

22/08/26

 

실버 랜덤 디펜스를 하다가 만난 문제이다.

 

두 문제가 시간제한과 N의 크기만 제외하고 완벽히 동일한 문제이므로 따로 문제를 분리하지 않고 합쳐서 글을 써본다.

 

이 글은 DP와 브루트 포스가 뭔지 대충 의미는 알고 있다는 전제하에서 쓰여있다.

 

---

 

이 문제는 두 가지 방법으로 접근 할 수 있는 좋은 문제이다.

 

첫 번째 풀이 방법은 DP이고, 두 번째 풀이 방법은 브루트 포스이다.

 

나는 처음 접근 했을 때 DP로 접근했었다.

 

특별한 이유는 없었고, 그냥 DP적인 풀이가 먼저 생각났기 때문이다.

 

이전에 캡틴 이다솜(1660번)문제를 푼 적이 있었는데, 그 문제와 접근 방식이 유사한 것 같아서 DP로 접근해보았다.

 

---

 

첫 번째 접근 방식:

 

 

DP적인 풀이 방식인데, 처음에 케이스를 두 가지로 나누었다.

 

---

1. N이 제곱수인가? -> 그렇다면 1 출력

 

2. N이 제곱수가 아닌가?

---

 

결국 우리는 N = (어떤 수) + (제곱 수)가 되도록 하고 싶다.

 

물론, 그 "어떤 수"도 제곱 수들의 합으로 이루어져 있을 것이다.

 

이때, 우리는 제곱 수들의 합으로 N을 표현하는데 그 제곱 수들의 개수가 최소가 되도록 해야 된다.

 

그렇다면 자연스럽게 dp[i] = i를 만들 때 필요로 하는 제곱 수들의 개수라고 하자.

 

그러면 우리는 dp[N]을 구하는 것으로 문제가 바뀌게 된다.

 

dp의 가장 큰 특징은 큰 문제를 작은 문제로 쪼개고, 그 작은 문제의 해답을 어딘가 저장해놓았다가 나중에 써먹는다는 점이다.

 

때문에 이 문제 또한 작은 문제로 분할하는 방법을 생각해야 할 것이다.

 

우리는 제곱 수들의 합으로 N을 표현한다고 했다.

 

그렇다면 다음과 같이 N이 표현될 것이다.

 

N = (제곱수들의 합) + (제곱 수)

 

 

우리는 (제곱수들의 합)이 최소가 되야지만 N이 최소가 된다는 사실을 알고 있다.

 

이 의문의 (제곱수들의 합)을 k라고 하자.(k는 여러 개가 될 수 있다.)

 

그렇다면, dp[N] = min(dp[k]) + 1이 된다는 사실을 알고 있다.

 

k를 식으로 표현하자면 N - (제곱수들)이므로, min(dp[N - (제곱수)]) + 1이 최종적인 식이 될 것이다.

 

이를 완성된 파이썬 코드로 작성한 결과는 다음과 같다. 더보기를 누르면 확인할 수 있다.

 

# 17626번 Four Squares
# DP
n = int(input())
square_num_li = [i*i for i in range(1, 224)]
dp = [0 for _ in range(n+1)]
for i in range(1, n+1):
    if i in square_num_li:
        dp[i] = 1
    else:
        dp[i] = min([dp[i-j] for j in square_num_li if i-j > 0])+1
print(dp[n])

# 1699번 제곱수의 합
# DP
n = int(input())
square_num_li = [i*i for i in range(1, 316)]
dp = [0 for _ in range(n+1)]
for i in range(1, n+1):
    if i in square_num_li:
        dp[i] = 1
    else:
        dp[i] = min([dp[i-j] for j in square_num_li if i-j > 0])+1
print(dp[n])

 

 

---

 

두 번째 접근 방식:

 

 

첫 번째 접근 방식은 17626번에서 시간 초과가 나게 만들었다.(python 3 제출의 경우에 한해서)

 

때문에 pypy로 제출하게 만들었는데, 평소 python3 제출로 어떻게든 맞고자 하는 나로서는 용납할 수 없는 제출이었다.

 

때문에 어떻게 접근을 할까 고민하던 중 태그에 브루트 포스가 있길래 그 방식대로 접근해보았다.

 

처음엔 어떻게 접근할지 몰라서, 나올 수 있는 모든 리스트를 미리 만들어놓고 그냥 출력만 하는 형식으로 제출했더니 맞았습니다를 받았다.

 

15만 바이트 짜리 코드인데, 아무래도 그 방식으로 푸는 것은 아닌 것 같아서 다른 사람의 코드를 참고하여 제출해보기로 하였다.

 

15만 바이트 짜리 코드는 보기 별로 좋은 코드는 아닌 것 같아 올려두지는 않겠다.

 

다른 분의 코드(kicza님)를 참고하여 작성해보았다. 

 

브루트 포스의 핵심 아이디어는 기껏 해봤자 최대 4개의 제곱수로 만들 수 있다는 점이었다.

 

---

1. 먼저 제곱 수면 당연히 1을 출력하면 된다.

 

2. 두 제곱수의 합으로 표현할 수 있다면 2를 출력하면 된다.(두 제곱수의 합을 저장하는 리스트를 따로 만들어준다,)

 

3. 제곱수 리스트를 반복하여 돌면서 N - (제곱수)를 해주는데, 이 수가 두 제곱수의 합 리스트에 있으면 3을 출력하면 된다. (N = (두 제곱수의 합) + (제곱수) 이므로)

 

4. 나머지 경우는 전부 4이다.

---

 

이 과정을 사용하여 코드를 작성하니 python3로도 충분히 돌아갔고, 코드 또한 직관적으로 보기가 좋았다.

 

아래는 이 과정을 거쳐 작성한 코드이다. 더보기를 누르면 확인할 수 있다.

 

# 17626번 Four Squares
# kicza님의 코드를 참고하여 작성하였습니다.
# 브루트 포스
from math import sqrt
from itertools import combinations_with_replacement

n = int(input())
square_num_li = [i*i for i in range(1, int(sqrt(n))+1)]
square_num_li_2 = [sum(k) for k in combinations_with_replacement(square_num_li, 2)]

def answer(n):
    if n in square_num_li: # 제곱수면
        return 1
    elif n in square_num_li_2: # 제곱수 두개를 더해서 만들 수 있는 수면
        return 2
    else:
        for square in square_num_li: # 제곱 수 중
            if n - square in square_num_li_2: # n에서 제곱수를 뺀 수가 제곱수 두개를 더해서 만들수 있는 수면
                return 3
    return 4

print(answer(n))

# 1699번 제곱수의 합
# 브루트 포스
from math import sqrt
from itertools import combinations_with_replacement

n = int(input())
square_num_li = [i*i for i in range(1, int(sqrt(n))+1)]
square_num_li_2 = [sum(k) for k in combinations_with_replacement(square_num_li, 2)]

def answer(n):
    if n in square_num_li:
        return 1
    elif n in square_num_li_2:
        return 2
    else:
        for square in square_num_li:
            if n - square in square_num_li_2:
                return 3
    return 4

print(answer(n))