지난 포스트에서 Two Sum을 구현했었는데 이번 Three Sum문제에서 활용할 수 있다.
먼저 문제는 아래와 같다.

Given an array nums of n integers, are there elements a, b, c in nums such that a + b + c = 0? Find all unique triplets in the array which gives the sum of zero.

Note:

The solution set must not contain duplicate triplets.

Example:

Given array nums = [-1, 0, 1, 2, -1, -4],

A solution set is:
[
  [-1, 0, 1],
  [-1, -1, 2]
]

위 예제처럼 배열 3개 요소를 더하여 0이 되는 요소들을 중복없이 반환하는 문제다.

이것을 의사코드로 작성해보면 다음과 같다.

0. Two Sum함수 작성
1. 배열요소를 오름차순 정렬한다.
2. 배열요소를 담을 해쉬셋를 생성
3. 입력받은 요소들이 2보다 클 동안 
    4. 배열요소에서 첫번째 요소를 추출한다.
    5. TwoSum(배열요소)함수에서 결과값을 받아서 리스트에 담는다.
    6. 리스트의 요소를 순회하면서
        7. 첫번째 요소가 리스트에 존재하면 해쉬셋에 추가한다.
8. 해쉬셋을 리스트로 반환한다.

코드로 구현하기 (파이썬)

0. Two Sum 함수 작성

먼저 지난 시간에 작성했던 Two Sum함수를 작성하면 아래와 같다.
원래는 처음 찾은 인덱스만 반환했었는데 일치하는 모든 데이터셋을 반환하도록 변경했다.

    def twoSum(self, nums, target):
        lookup = {}
        hashset = set()
        for index, item in enumerate(nums):
            if item in lookup : hashset.add((nums[lookup[item]], nums[index]))
            else : lookup[target - item] = index
        return hashset

1-8. 구현하기

    def threeSum(self, nums):
        nums.sort()
        res = set()
        while len(nums) > 2 :
            key = nums.pop(0)
            tmp = self.twoSum(nums, -key)
            for item in tmp:
                res.add((key, item[0], item[1]))
        return [list(item) for item in res]

시간복잡도는: O(n^2)인 알고리즘이며 효율성보다는 심플한 알고리즘 설계를 목표로 했다.
실행시켜보면 원하는 결과 값을 확인 할 수 있는데 LeetCode에서 313개의 코드테스트 결과 1.6초가 나왔다.
다음번 포스트에서는 같은 문제를 효율성을 높인 다른 알고리즘으로 풀어 볼 것이다.

solution = Solution()
print(solution.threeSum([-4,-2,-2,-2,0,1,2,2,2,3,3,4,4,6,6]))

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

Runtime: 1604 ms, faster than 32.52% of Python3 online submissions for 3Sum.

Three Sum 스피드개선 보러가기