[Leetcode] 15. 3Sum

문제 바로가기

Time Complexity: O(n^2)
Space Complexity: O(n)

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        n = len(nums)
        nums.sort()
        answer = list()
        
        for i in range(n-2):
            target = nums[i]
            if target > 0:
                break
            if i == 0 or nums[i-1] != nums[i]:
                l = i + 1
                r = n - 1
                while l < r:
                    if -target == nums[l] + nums[r]:
                        answer.append([target, nums[l], nums[r]])
                        # 중복 제거
                        while l < r and nums[l] == nums[l+1]:
                            l += 1
                        while l < r and nums[r] == nums[r-1]:
                            r -= 1
                        l += 1    
                        r -= 1
                    elif nums[l] + nums[r] > -target:
                        r -= 1
                    else:
                        l += 1
        
        return answer

Counter 활용

Using collection.Counter get all unique numbers with count of occurrences.

from collections import Counter

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        res = []
        nums = Counter(nums)
        keys = sorted(nums.keys(), reverse=True)
        while keys:
            keysLen = len(keys)
            n1 = keys[keysLen - 1]
            if n1 > 0:
                break
            for n2 in keys[(keysLen - 1 if nums[n1] > 1 else keysLen - 2):]:
                if n2 > -n1:
                    break
                n3 = -(n1+n2)
                if (n3 > n2 and n3 in nums) or (n3 == n2 and nums[n2] > 1 + (1 if n2 == n1 else 0)):
                    res.append([n1,n2,n3])
            keys.pop()
        return res