[LeetCode] 15. 3Sum

Problem:

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]]

Example 2:
Input: nums = []
Output: []

Example 3:
Input: nums = [0]
Output: []

Solution:

세 개의 수를 더해서 0이 되는 것은 크게 3가지로 나뉠 수 있다.

세 개의 수가 동일할 때: 세 개의 수가 모두 0인 경우밖에 없다.
두 개의 수가 동일하고 하나는 다를 때
세 개의 수가 모두 다를 때

두 값이 주어졌을 때 나머지 하나의 값은 두 값의 합의 마이너스로 결정된다. 이 값이 존재하는지 안하는지 체크하는 시간을 줄이기 위해선 dictionary를 이용하여 탐색을 하면 $O(1)$의 시간복잡도로 해결할 수 있다. 위 두 가지 사실을 이용하여 구현하면 문제를 빠르게 해결할 수 있다.

Code:

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        if len(nums) < 3:
            return []
        
        counts = {}
        for num in nums:
            counts[num] = counts[num] + 1 if num in counts else 1
        
        cases = []
        if counts.get(0):
            if counts[0] >= 3:
                cases = [[0,0,0]]
                
        nums = list(sorted(counts.keys()))
        
        for i in range(len(nums)):
            if counts.get(nums[i]) > 1 and nums[i] != 0:
                if counts.get(-nums[i] * 2):
                    cases.append([nums[i], nums[i], -2 * nums[i]])
            for j in range(i+1, len(nums)):
                if counts.get(-nums[i]-nums[j]):
                    if nums[j] >= -nums[i]-nums[j]:
                        break
                    else:
                        cases.append([nums[i], nums[j], -nums[i]-nums[j]])
        return cases

← Previous Post Next Post →