Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

파이썬 백트래킹으로 인접한 쌍의 합이 완전제곱수가 되는 순열 개수 구하기

숫자 리스트 nums가 주어졌을 때, 인접한 모든 값의 쌍의 합이 완전제곱수(perfect square)가 되도록 만드는 순열의 개수를 구하는 문제입니다. 두 순열 A와 B는 임의의 인덱스 i에서 A[i]와 B[i]가 서로 다른 경우가 하나라도 존재하면 서로 다른 순열로 간주합니다.

예를 들어 입력이 nums = [2, 9, 7]이라면 출력은 2입니다. 조건을 만족하는 순열이 [2, 7, 9][9, 7, 2], 이렇게 두 가지이기 때문입니다.

알고리즘 설계

이 문제는 백트래킹(backtracking) 기법으로 효율적으로 해결할 수 있습니다. 각 위치에 올 수 있는 후보를 하나씩 배치해 보고, 조건을 만족하지 않으면 이전 상태로 되돌아가 다른 후보를 시도하는 방식입니다. 구체적인 절차는 다음과 같습니다.

  1. 결과를 저장할 변수 res를 0으로 초기화합니다.
  2. 재귀 함수 util(i)를 정의합니다.
  3. i + 1이 nums의 길이와 같다면 마지막 위치까지 유효한 배치를 완성한 것이므로 res를 1 증가시키고 함수를 종료합니다.
  4. 현재 단계에서 이미 시도한 합을 기록할 빈 집합 visited를 생성합니다.
  5. j를 i + 1부터 nums의 끝까지 반복하며 다음을 수행합니다.
    • s = nums[i] + nums[j]를 계산합니다.
    • s가 아직 시도되지 않았고, (√s)² == s, 즉 s가 완전제곱수라면:
      • s를 visited에 추가합니다.
      • nums[i + 1]과 nums[j]의 자리를 맞바꿉니다.
      • util(i + 1)을 재귀 호출합니다.
      • 자리 교환을 되돌려 원래 상태로 복원합니다(백트래킹).
  6. 메인 함수에서는 다음을 수행합니다.
    • 새로운 집합 visited를 생성합니다.
    • i를 0부터 nums의 길이까지 반복하며:
      • nums[i]와 nums[0]의 자리를 맞바꿉니다.
      • nums[0]이 아직 시도되지 않은 값이라면 util(0)을 호출합니다.
      • nums[0]을 visited에 추가합니다.
      • 자리 교환을 되돌립니다.
  7. res를 반환합니다.

visited 집합을 사용하는 핵심 이유는 중복 제거입니다. 리스트에 같은 값이 여러 개 있을 때, 동일한 값이 같은 위치에 놓이는 분기는 결과가 항상 같으므로 한 번만 탐색하면 됩니다. 이를 통해 불필요한 중복 계산을 크게 줄일 수 있습니다.

구현 예제

from math import sqrt
class Solution:
    def solve(self, nums):
        self.res = 0
        def util(i):
            if i + 1 == len(nums):
                self.res += 1
                return
            visited = set()
            for j in range(i + 1, len(nums)):
                s = nums[i] + nums[j]
                if s not in visited and int(sqrt(s)) ** 2 == s:
                    visited.add(s)
                    nums[i + 1], nums[j] = nums[j], nums[i + 1]
                    util(i + 1)
                    nums[i + 1], nums[j] = nums[j], nums[i + 1]
        visited = set()
        for i in range(len(nums)):
            nums[i], nums[0] = nums[0], nums[i]
            if nums[0] not in visited:
                util(0)
            visited.add(nums[0])
            nums[i], nums[0] = nums[0], nums[i]
        return self.res
ob = Solution()
nums = [2, 9, 7]
print(ob.solve(nums))

입력

[2, 9, 7]

출력

2

마무리 정리

이 풀이는 백트래킹으로 가능한 모든 배치를 탐색하면서, 완전제곱수 조건을 만족하지 않는 가지는 조기에 잘라내는 가지치기(pruning)를 적용하고, visited 집합으로 중복 값에 의한 동일 순열의 재탐색까지 방지합니다. 그 결과 단순한 전수조사보다 훨씬 효율적으로 정답을 구할 수 있습니다.