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

파이썬으로 풀는 3Sum 문제: 합이 0이 되는 세 숫자 조합 찾기

숫자 배열이 주어졌을 때, 배열에 저장된 n개의 정수 중에서 세 원소 a, b, c를 골라 그 합이 a + b + c = 0이 되는 모든 고유한 조합(triplet)을 찾는 것이 이번 문제의 목표입니다.

예를 들어 배열이 [-1, 0, 1, 2, -1, -4]와 같다면, 결과는 [[-1, -1, 2], [-1, 0, 1]]이 됩니다.

해결 접근 방법

이 문제는 정렬과 두 포인터(Two Pointers) 기법을 활용하면 O(n²) 시간 복잡도로 효율적으로 해결할 수 있습니다. 단계별로 살펴보겠습니다.

  1. 배열 nums를 오름차순으로 정렬하고, 결과를 저장할 빈 배열 res를 선언합니다.
  2. i를 0부터 (배열 길이 - 3)까지 반복합니다.
    • i > 0이면서 nums[i]가 nums[i-1]과 같다면 중복 조합이 되므로 건너뜁니다.
    • 왼쪽 포인터 l := i + 1, 오른쪽 포인터 r := 배열 길이 - 1로 초기화합니다.
    • l < r인 동안 다음을 반복합니다:
      • sum := nums[i] + nums[l] + nums[r]
      • sum < 0이면 l을 1 증가시키고, sum > 0이면 r을 1 감소시킵니다.
      • sum == 0이라면 [nums[i], nums[l], nums[r]]을 res에 추가합니다.
      • 중복 값을 건너뛰기 위해 l은 연속된 같은 값이 끝날 때까지 증가시키고, r은 같은 값이 끝날 때까지 감소시킵니다.
      • 마지막으로 l을 1 증가시키고 r을 1 감소시켜 새로운 조합을 탐색합니다.
  3. 모든 탐색이 끝나면 res를 반환합니다.

파이썬 구현 예제

아래 코드를 통해 실제 동작 방식을 더 잘 이해할 수 있습니다.

class Solution(object):
    def threeSum(self, nums):
        nums.sort()
        result = []
        for i in range(len(nums)-2):
            if i > 0 and nums[i] == nums[i-1]:
                continue
            l = i+1
            r = len(nums)-1
            while(l<r):
                sum = nums[i] + nums[l] + nums[r]
                if sum<0:
                    l+=1
                elif sum >0:
                    r-=1
                else:
                    result.append([nums[i],nums[l],nums[r]])
                    while l<len(nums)-1 and nums[l] == nums[l + 1] : l += 1
                    while r>0 and nums[r] == nums[r - 1]: r -= 1
                    l+=1
                    r-=1
        return result
ob1 = Solution()
print(ob1.threeSum([-1,0,1,2,-1,-4]))

입력

[-1,0,1,2,-1,-4]

출력

[[-1,-1,2],[-1,0,1]]

핵심 포인트 정리

  • 정렬 우선: 배열을 먼저 정렬해야 두 포인터 기법을 적용할 수 있고, 중복 제거도 쉬워집니다.
  • 두 포인터 활용: 합이 음수면 왼쪽 포인터를, 양수면 오른쪽 포인터를 이동시켜 0에 가까워지도록 탐색 범위를 좁힙니다.
  • 중복 처리: 같은 값이 연속될 경우 포인터를 추가로 이동시켜 동일한 조합이 여러 번 출력되지 않도록 합니다.
  • 시간 복잡도: 정렬에 O(n log n), 전체 탐색에 O(n²)이 소요되므로 총 O(n²)입니다.