[파이썬 알고리즘 인터뷰] 9. 세 수의 합

github.com/onlybooks/algorithm-interview

1. 문제

15. 3Sum

 nums 배열을 입력받으면 서로 다른 i, j, k에 대해 nums[i] + nums[j] +nums[k] == 0인 [nums[i], nums[j], nums[k]] 배열을 리턴하시오. 다만, 같은 세 수가 여러번 포함되면 안된다.

 

2. 딕셔너리를 이용한 내 풀이

 리트코드 1번 문제 두 수의 합에서 사용했던 딕셔너리를 이용했다. 먼저 입력받은 배열을 숫자를 key, 개수를 value로 하는 딕셔너리를 만들어 저장한다. 그리고 딕셔너리의 모든 key에서 첫번째 수, 두번째 수를 찾는다. 중복을 막기 위해서 사용한 수는 value에서 1을 빼고, 첫번째 수는 두번째 수 이상, 두번째 수는 세번째 수 이상이여야 한다고 조건을 건다. 두번째 수까지 찾으면 첫번째 수 + 두번째 수 + 세번째 수 = 0이므로 첫번째 수 + 두번째 수 = -세번째 수인 성질을 이용해 -(첫번째 수 + 두번째 수)를 key로 딕셔너리에서 세번째 수가 존재하는지 확인한다. 세번째 수가 존재한다면 이 숫자들을 배열에 추가한다. O(n^2)의 시간복잡도를 가지고, 실행시간은 2500ms 정도를 보였다.

def threeSum(self, nums: List[int]) -> List[List[int]]:
    result = []
    num_dict = collections.defaultdict(int)
    for n in nums:
        num_dict[n] += 1

    for first in num_dict.keys():
        num_dict[first] -= 1
        for second in num_dict.keys():
            if num_dict[second] < 1 or second > first:
                continue
            num_dict[second] -= 1
            if -(first + second) in num_dict.keys() and num_dict[-(first + second)] > 0 and second >= -(first + second):
                result.append([first, second, -(first + second)])
            num_dict[second] += 1
        num_dict[first] += 1

    return result

3. 투 포인터를 이용한 풀이

 이 방법은 앞선 문제에서도 자주 등장했던 투 포인터를 이용한 방법이다. 먼저 같은 수 처리를 쉽게 하기 위해서 배열을 정렬한다. 가장 오른쪽에 있는 수를 right, i 뒤에 수 중에서 가장 왼쪽에 있는 수를 left로 정한다. 그리고 nums[i] + nums[right] + nums[left]의 합이 0 이상이면 right를 왼쪽으로, 0 이하면 left를 오른쪽으로 한칸 옮긴다. 만약 세 수의 합이 0이라면 정답 배열에 값을 추가하고 left는 한칸 오른쪽으로, right는 한칸 왼쪽으로 옮긴다. 만약 right, left를 옮겼는데 같은 값이라면 한칸 더 옮긴다.  이 과정을 리스트 안에 모든 i에 대해서 진행한다. 마찬가지로 O(n^2)의 시간복잡도를 가지지만 880ms의 실행속도로 2번보다 3배 정도 빨리 실행되었다.

def threeSum(self, nums: List[int]) -> List[List[int]]:
    results = []
    nums.sort()

    for i in range(len(nums) - 2):
        # 중복된 값 건너뛰기
        if i > 0 and nums[i] == nums[i - 1]:
            continue

        # 간격을 좁혀가며 합 `sum` 계산
        left, right = i + 1, len(nums) - 1
        while left < right:
            sum = nums[i] + nums[left] + nums[right]
            if sum < 0:
                left += 1
            elif sum > 0:
                right -= 1
            else:
                # `sum = 0`인 경우이므로 정답 및 스킵 처리
                results.append([nums[i], nums[left], nums[right]])

                while left < right and nums[left] == nums[left + 1]:
                    left += 1
                while left < right and nums[right] == nums[right - 1]:
                    right -= 1
                left += 1
                right -= 1

    return results

 이 문제는 너무 효율적인 풀이를 찾으려다가 제대로 못찾아서 문제 푸는 걸 포기할 뻔했다. 문제를 아예 못푸는 것보다는 약간 비효율적 풀이라도 일단 작동하는 풀이를 찾는 게 좋은 것 같다.