백준 1699 제곱수의 합 - python

# 중복조합
from itertools import combinations_with_replacement

N = int(input())
a = [int(n**2) for n in range(1, int(N**0.5) + 1)]

if int(N**0.5)**2 == N:
    print(1)
else:
    i = 2
    while True:
        tag = 0
        if a[0]*i <= N <= a[-1]*i:
            for x in combinations_with_replacement(a, i):
                if sum(x) == N:
                    tag += 1
                    break
            if tag>0:
                print(i)
                break
            else:
                i+=1
        else:
            i+=1

N보다 작은 모든 제곱수를 넣은 리스트 a를 만들었습니다.

그리고 리스트 a에서 중복조합으로 i 개를 뽑아서 합을 구한 뒤 N과 같은 값을 갖는지 확인했습니다.

예) N = 17, a = [1, 4, 9, 16], i = 2라면

(1, 1), (1, 4), (1, 9), (1, 16), (4, 4), (4, 9), (4, 16), (9, 9), (9, 16), (16, 16)

이 뽑힐 것이고 (1, 1)의 합인 2부터 for 문을 돌면서 확인한 뒤, (1, 16)의 합인 17이 N과 같으므로 i 값인 2를 답으로 출력하게 됩니다. (만약 답이 없다면 i에 1을 더해서 i = 3인 경우로 다시 돕니다.)

i = 1인 경우는 N이 제곱수인지 아닌지 확인하는 방법으로 대체했습니다.

itertools를 통해 순열, 중복순열, 조합, 중복조합을 이용할 수 있습니다.

# 순열(순서 O, 중복 X)
from itertools import permutations

# 중복순열(순서 O, 중복 O)
from itertools import product

# 조합(순서 X, 중복 X)
from itertools import combinations

# 중복조합(순서 X, 중복 O)
from itertools import combinations_with_replacement