[리트코드] 15. 3Sum

리트코드 15번 문제.

난이도: Medium

 

문제

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.

 

예시

Example 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.

Example 2:

Input: nums = [0,1,1]
Output: []
Explanation: The only possible triplet does not sum up to 0.

Example 3:

Input: nums = [0,0,0]
Output: [[0,0,0]]
Explanation: The only possible triplet sums up to 0.

 

답

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        res = []
        nums.sort() #1
        
        for left in range(len(nums)- 2): #2
            if left != 0 and nums[left] == nums[left - 1]: #3
                continue
            
            mid = left + 1 #4
            right = len(nums) - 1
            while mid < right: #5
                three_sum = nums[left] + nums[mid] + nums[right] #6
                if three_sum < 0:
                    mid += 1
                elif three_sum > 0:
                    right -= 1
                else:
                    res.append([nums[left], nums[mid], nums[right]]) #7
                    
                    while mid < right and nums[mid] == nums[mid + 1]: #8
                        mid += 1
                    
                    while mid < right and nums[right] == nums[right - 1]: #9
                        right -= 1
                    
                    mid += 1 #10
                    right -= 1
            
        return res

Time Complexity: O(n^2)

 

 

설명

이번 문제는 주어진 정수 리스트에서 합이 0이 되는 세 쌍의 수를 찾아내 리턴해야 하는 문제다. 이 문제의 키 포인트는 리스트를 먼저 sort 하는 것이다. 오름차순으로 정렬이 된 리스트를 이용하면 그렇지 않은 리스트와는 다르게 여러 트릭을 쓸 수 있다. 참고로 sort의 Time Complexity는 O(nlogn).

 

정렬이 된 리스트에서 left, mid, right 세 변수가 0이 되는 index 세 쌍을 찾아낼 것이다. left는 가장 작은 수의 index, mid는 두 번째, right는 가장 큰 수의 index가 되는 방식이다. 여러 루프를 이용해 세 쌍이 될 수 있는 모든 경우를 체크한다.

 

#2 블록의 for 루프는 len(nums) - 2까지만 돌린다. 예를 들어 len(nums) = 7인 리스트에서 len(nums) - 2 = 5, for 루프는 0부터 4까지 돌아간다. left가 5나 6일 때는 체크할 필요가 없다. mid와 high가 되는 두 인덱스가 필요하기 때문이다.

 

#3 블록에서는 left와 left - 1이 중복될 경우를 체크한다. 문제 description에서 나와있듯이 중복된 쌍은 리턴하면 안 된다. 예를 들어 [-1, -1, 0, 1]인 리스트에서 -1 + 0 + 1 = 0은 첫 번째 -1와 두 번째 -1로 만들 수 있기 때문에 leftmost인 left를 쓰는 방식이다. 



#4에서 mid와 right을 설정한다. left보다 큰 인덱스들을 이용해서 쌍을 찾아야 하기 때문에 mid는 그중 가장 작은 수, right은 그중 가장 큰 수로 설정한다.

 

#5에서 while루프를 돌린다. 이 루프에서 설정된 left, mid, right 쌍의 합을 계산하고 합이 0보다 적으면 mid를 1 올리고 합이 0보다 크면 right을 내린다. 이것이 #1에서 sort를 한 이유다. 더 큰 합을 만들기 위해 mid와 right을 조절할 수 있기 때문이다.

 

#7에서 합이 0인 쌍을 찾았으면 그 쌍을 리턴 밸류에 넣는다. #8과 #9에서 또 while loop이 나온다. #3과 비슷하게 이번에는 중복되는 쌍을 제거하기 위해 사용된다. 마지막으로 #10에서 mid와 right이 다시 설정된다.

 

#4의 while loop이 끝나면 마지막으로 찾은 세 쌍이 기록된 res를 리턴한다.

 

노트

이 문제 또한 회사 인터뷰에 수도 없이 나온다. 보통 인턴이나 entry-level 엔지니어 인터뷰에 많이 나오는 듯하다. 개인적으로 백여 번의 technical interviews를 봤는데 두세 번은 이 문제가 나온 것 같다. 그만큼 유명하기도 하고 sorting, pointer 등 여러 콘셉트도 나오니 잘 짜인 문제이다. 개발자가 되려면 무조건 풀어봐야 하는 문제다.

 

여담이지만 이 문제의 logic이 메이플스토리 파티 퀘스트와 좀 닮은 것 같다. 유저가 1부터 9까지 구성된 숫자 상자에 서서 임의로 지정된 발판 다섯 가지를 찾아야 하는 퀘스트인데 12345부터 12346, 12347,... , 56789까지 모든 경우의 수를 시도해 봐야 한다. 경우의 수는 9C5이니 126.

이미지 1: 루디브리엄 파티퀘스트 발판서기

 

 

출처

이미지 1 출처: https://namu.wiki/w/%ED%8C%8C%ED%8B%B0%20%ED%80%98%EC%8A%A4%ED%8A%B8