파이썬으로 서로 다른 숫자의 n제곱 합이 특정 수가 되는 조합의 수 구하기

숫자 x와 n이 주어졌을 때, 서로 중복되지 않는 고유한 숫자들의 n제곱의 합으로 x를 표현할 수 있는 방법의 수를 구하는 문제입니다.

예를 들어 입력이 x = 100, n = 2라면 출력은 3이 됩니다. 가능한 조합은 다음과 같기 때문입니다.

• 6² + 8²
• 10²
• 1² + 3² + 4² + 5² + 7²

알고리즘 접근 방법

이 문제는 재귀 호출을 활용한 백트래킹 기법으로 해결할 수 있습니다. 현재까지 선택한 숫자들의 거듭제곱 합(cs)과 다음 후보 숫자(cn)를 추적하면서, 목표값 x에 정확히 도달하는 경우의 수를 하나씩 세어 나갑니다.

1. 정답을 저장할 변수 ans를 0으로 초기화합니다.
2. x, n, cn, cs 네 개의 매개변수를 받는 solve() 메서드를 정의합니다. 초기값은 cs = 0, cn = 1입니다.
3. p를 cn의 n제곱 값으로 설정합니다.
4. p + cs가 x보다 작은 동안 다음 과정을 반복합니다.
   • ans에 solve(x, n, cn + 1, p + cs)의 반환값을 더합니다.
   • cn을 1 증가시킨 뒤, p를 새로운 cn의 n제곱 값으로 갱신합니다.
5. 반복이 끝난 후 p + cs가 x와 같다면 ans를 1 증가시킵니다.
6. ans를 반환합니다.

구현 예시

아래 파이썬 코드를 통해 동작 방식을 더 자세히 살펴보겠습니다.

from math import pow

def solve(x, n, cn = 1, cs = 0):
    ans = 0
    p = pow(cn, n)
    while p + cs < x:
        ans += solve(x, n, cn + 1, p + cs)
        cn = cn + 1
        p = pow(cn, n)

    if p + cs == x:
        ans = ans + 1
    return ans

x = 100
n = 2
print(solve(x, n))

입력

100, 2

출력

3

이 코드는 1부터 시작해 각 숫자를 현재 합에 포함할지 여부를 재귀적으로 결정합니다. 이미 목표값을 초과한 경로는 더 이상 탐색하지 않으므로 불필요한 연산을 줄일 수 있으며, 최종적으로 유효한 조합의 총 개수를 반환하게 됩니다.