3Sum (in Python)

Problem

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

Solution

•배열 정렬.

•i를 기준으로 left boundary와 right boundary를 설정.

1. 세 수의 합이 0보다 작을 경우 left boundary를 오른쪽으로 이동.

2. 세 수의 합이 0보다 클 경우 right boundary를 왼쪽으로 이동.

-> 정렬 되어있는 배열이기 때문.

•a+b+c=0을 만족하는 세 수를 발견하면 ans 배열에 [a,b,c] 배열을 원소로 삽입.

class Solution:
    def threeSum(self, nums):
        """
        :type nums: List[int]
        :rtype: List[int]
        """
 
        ans, n = [], len(nums)        
        nums.sort()
 
        for i in range(n) :
 
            if i > 0 and nums[i] == nums[i-1] :
                continue
 
            # 범위 지정해서 탐색
            l, r = i+1, n-1
 
            while l < r :
                s = nums[i] + nums[l] + nums[r]
                if s < 0 :
                    l += 1
                elif s > 0 :
                    r -= 1
                else :
                    # 탐색 완료
                    ans.append([nums[l],nums[i],nums[r]])
                    while l < n-1 and nums[l] == nums[l+1] :
                        l += 1
                    while r > 1 and nums[r] == nums[r-1] :
                        r -= 1
                    l += 1
                    r -= 1
 
        return ans

Complexity Analysis

•Time complexity : O(n^2)

•Space complexity : O(1)

Source :

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