백준 17626 python

문제 링크

몇일 전에 코딩테스트 스터디에서 어떤 문제였는지 기억이 나지 않지만 combination을 쓰길레 그게 기억이 남아 반복이 있는 combination을 써서 문제를 풀어봤다.

사실 Cubes 라는 문제를 계속 트라이중인데 처음에는 이걸 생각하면서 풀다보니 처음은 풀이를 좀 이상하게 해서 틀리고, 그 다음은 3에서도 없었을 때 4는 굳이 계산할 필요가 없는데 계산해서 시간초과가 나더라..

from itertools import combinations_with_replacement as cr

num = int(input())
tmp = 1
sqr = []
while tmp**2<=num:
    sqr.append(tmp**2)
    tmp += 1

if num == sqr[-1]:
    print(1)
    exit(0)

for i in range(2,4):
    for j in cr(sqr,i):
        if sum(j) == num:
            print(i)
            exit(0)

print(4)  # 2도 아니고 3도 아니면 갯수는 무조껀 4겠지..

또 풀고나서 다른 사람들의 시간을 구경해봤다가 더 효율적인 방법이 있다는 것을 알고 그거를 봤는데 그냥 단순히 계산하는 것이라는 것을 알아가지고 다르게 풀어본 것도 올려본다.

def check(num):
    sqr = num**(0.5)
    if sqr.is_integer():
        print(1)
        return 
    
    for i in range(1,int(sqr)+1):
        if ((num-i*i)**(0.5)).is_integer():
            print(2)
            return
    
    for i in range(1,int(sqr)+1):
        for j in range(i,int(sqr)+1):
                res = (num - i*i - j*j)
                if res<=0:
                    continue
                
                if (res**(0.5)).is_integer():
                    print(3)
                    return
    
    print(4)
    return
    
check(int(input()))