[Baekjoon]백준 17626 Four Squares(실버 3) - Python

문제설명

예제

문제를 살펴보면 모든 자연수가 4개 이하의 제곱수의 합으로 표현가능하다고한다

주어진 수를 몇개의 제곱수 합으로 표현가능한지 구하는 문제이다

 

문제를 보자말자 가장 먼저 든 생각은 제곱수를 활용해서 위에서부터 내려가며 제곱수로 표현가능한지 확인하는것이었다

이전에 풀었던 문제에서 sol()함수를 구현했어서 해당 함수를 수정하여 코드를 완성하였다

먼저 if문을 통해 현재 입력받은 수의 제곱근을 구하였을때, 제곱근이 정수인지 확인해주었다

제곱근이 정수가 아니라면 하나의 제곱수로 나타낼 수 없다는 뜻이 된다

이후 for문을 통해 특정 제곱수를 입력받은 k에서 뺀 후에 해당 값(제곱근)이 정수인지 확인하면된다

3개의 제곱수가 필요한 경우는 for문을 두 개 만들어 확인해주면 된다

최종적으로 3개의 제곱수로도 표현이 되지 않는다면 무조건 4개의 제곱수로 표현되므로(문제 조건)

return 4를 하면서 함수를 종료하면된다

import sys
input = sys.stdin.readline
import math as m

def sol(k):
    if int(m.sqrt(k)) == m.sqrt(k):
        return 1
    for i in range(1, int(m.sqrt(k))+1):
        if int(m.sqrt(k-i**2)) == m.sqrt(k-i**2):
            return 2
    for i in range(1, int(m.sqrt(k))+1):
        for j in range(1, int(m.sqrt(k-i**2))+1):
            if int(m.sqrt(k-i**2-j**2)) == m.sqrt(k-i**2-j**2):
                return 3
    return 4

num = int(input())
print(sol(num))