[완전 탐색] 백준 17626번 (파이썬) : Four Squares

문제 해석을 하는 것은 어렵지 않았습니다. 요약하가면 아무 숫자나 입력받고 그 숫자를 최대 4개 이하의 제곱수의 합으로 나타내기만 하면 된다는 것입니다.

import math

def four_squares(n):

  if int(math.sqrt(n)) == math.sqrt(n):
    return 1

  for i in range(1, int(math.sqrt(n)) + 1):

    if int(math.sqrt(n - pow(i, 2))) == math.sqrt(n - pow(i, 2)):
      return 2

  for i in range(1, int(math.sqrt(n)) + 1):
        for j in range(1, int(math.sqrt(n - pow(i,2))) + 1):
            if int(math.sqrt(n - pow(i,2) - pow(j,2))) == math.sqrt(n - pow(i,2) - pow(j,2)):
                return 3

  return 4

n = int(input())
print(four_squares(n))

구글링 해본 결과 DP나 브루트포스는 모르겠고 생각나는대로 제어문으로 도배했습니다. math에 있는 sqrt를 사용합니다. sqrt는 제곱근입니다. 

먼저 출력이 1인 경우입니다. 출력이 1이 되기 위해서는 입력한 숫자 자체가 제곱수여야 합니다. 예를 들어 25를 입력하면 5x5라서 1이기 때문입니다. 함수의 첫 줄의 if문이 이를 보여줍니다. 거창하게 ==으로 적혀있지만 속뜻은 제곱근이 정수로 바뀐 경우나 실수형인 경우나 같냐는 것인데 당연히 입력값 자체가 재곱수라면 같을 것입니다. 그러므로 반환값을 1입니다

 

다음은 for문이 등장합니다. range 내에 값은 1부터 n의 제곱근 + 1까지의 숫자이기 때문에 범위를 위와 같이 설정합니다. 그리고 if문에서는 입력값 n에서 1**2부터 sqrt(n)**2까지를 빼면서 이를 뺀 나머지 값이 제곱수 인지를 따지는 것입니다. 만약 조건문에 걸린다면 이는 두 가지 제곱수의 합으로 이루어져 있다고 할 수 있기 때문에 반환값이 2입니다.

 

마지막 for문중 바깥 for문은 앞선 return 2와 범위가 같습니다. 그러나 이제는 제곱수 3개로 이루어진 경우를 알아봐야하기 때문에 제곱수를 하나 더 빼줘야 합니다. 그래서 이중 for문으로 안에는 하나 더 빼줄 j**2를 넣어주는 것입니다. 원리는 위와 비슷하기 때문에 생략하겠습니다.

 

return 4같은 경우에는 조건/반복문을 만들 수도 있지만 문제에서 어차피 모든 숫자는 최대 4개의 제곱수의 합으로 이루어진다고 했기 때문에 앞선 모든 과정에도 들어가지 않는다면 그냥 4입니다. 그래서 시간을 줄이기 위해 return 4 하나로 마무리하였습니다.

 

(참고로 pow() 함수는 파이썬에 내장되어있는 함수로 지수함수를 뜻합니다.)