[LeetCode] 13. 3Sum

https://leetcode.com/problems/3sum/

 

LeetCode 13번 3Sum 문제는 다음과 같다.

 

Given an integer array nums, return all the triplets [nums[i], nums[j], nums[k]] such that i != j, i != k, and j != k, and nums[i] + nums[j] + nums[k] == 0.

Notice that the solution set must not contain duplicate triplets.

Constraints:
3 <= nums.length <= 3000
-10^5 <= nums[i] <= 10^5

 

 주어진 정수가 담긴 배열에서 세 개의 요소의 합이 0이 되는 조합을 반환하면 된다. 정답엔 중복된 조합은 없어야 된다.

세 개의 예시가 제시돼 있지만 여기선 주어진 예시 중 첫 번째 예시만 인용하겠다.

 

예시 1:

Input: nums = [-1,0,1,2,-1,-4]
Output: [[-1,-1,2],[-1,0,1]]
Explanation: 
nums[0] + nums[1] + nums[2] = (-1) + 0 + 1 = 0.
nums[1] + nums[2] + nums[4] = 0 + 1 + (-1) = 0.
nums[0] + nums[3] + nums[4] = (-1) + 2 + (-1) = 0.
The distinct triplets are [-1,0,1] and [-1,-1,2].
Notice that the order of the output and the order of the triplets does not matter.

 

 최적의 답안이 생각나면 좋겠지만 그렇지 못할 땐 brute force로 구현을 해보는 것도 좋은 방법이다.

아래는 위 문제에 대해 구현한 brute force로 input에 따른 올바른 output이 나오지만 Time Limit Exceeded에 걸리게 된다.

for i in range(n):
    for j in range(i+1, n):
        for k in range(j+1, n):
            if nums[i] + nums[j] + nums[k] == 0:
                arr = arr + [[nums[i], nums[j], nums[k]]]

 이 문제에는 두 가지 방법으로 정답을 구할 수 있다. 하나는 두 개의 포인터를 사용하는 방법이고, 다른 하나는 해시를 사용하는 방법이다. 여기선 두 개의 포인터를 사용하는 방법으로 구한 정답을 공개하겠다.

 두 개의 포인터를 이용하는 방법은 우선 주어진 nums 배열을 정렬을 하는 것으로 시작한다. 이후 구해야하는 세 개의 요소 중 첫 번째 요소 하나를 정렬된 배열에 처음부터 순서대로(linear) 뽑고 구해야하는 나머지 두 요소 중 두 번째 요소는 첫 번째 요소의 바로 다음으로 잡고, 마지막 세 번쨰 요소는 배열에 가장 마지막으로 포인터를 정해 구하는 방법이다.

 이 문제에서 내가 가장 어려웠던 부분은 중복을 없애는 것이었다. 중복을 세지 않는 아이디어는 뒤에 같은 정수가 따라오고 있다면 해당 정수는 건너뛰는 것이다.

 아래 정답에 target은 문제에서 원하는 합에 따라 다르게 설정할 수 있어 세 요소의 합이 0이나 7 등 약간의 변형에 얼마든지 적용이 가능하다.

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        i = 0
        j = i + 1
        k = len(nums) - 1
        target = 0
        ans = list()        
        nums.sort()
        
        for i in range(k-1):
            if i != 0 and nums[i-1] == nums[i]:
                continue
            j = i + 1
            k = len(nums) - 1
            while j < k:
                tot = nums[i] + nums[j] + nums[k]
                if tot < target:
                    j += 1                    
                elif tot > target:
                    k -= 1
                else:
                    ans += [[nums[i], nums[j], nums[k]]]
                    k -= 1
                    for j in range(j, k):
                        j += 1
                        if nums[j] != nums[j-1]:
                            break
        return ans

 두 개의 포인터를 이용하는 정답은 13번 3Sum 문제를 풀어보기 전에 167. Two Sum II - Input Array Is Sorted 문제를 먼저 풀어보는 것을 추천한다. 리트코드의 Solution에 따르면 면접관이 Two Sum 문제를 먼저 물어본 후 3Sum 문제를 제시하는 방식으로 면접이 이뤄질 수 있다고 한다. 지금까지 몇 번의 면접을 봤지만 단계적으로 문제를 제시한 적은 없지만 얼마든지 그럴 수 있고, 3Sum 문제에 대한 정답이 떠오르지 않는다면 3Sum 문제 정답의 기본 원리가 적용된 Two Sum 문제를 먼저 풀어본 후 아이디어를 얻어 3Sum 문제를 풀어보는 것을 추천한다.

https://leetcode.com/problems/two-sum-ii-input-array-is-sorted/solution/