LeetCode - 15. 3Sum [Python]

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

3Sum - LeetCode

---

풀이 과정

배열에서 합이 0을 이루는 3개의 원소를 출력하는 문제이다.

 

가장 간단하게 생각하면 완전 탐색으로 O(N^3)으로 풀 수 있을 것이다. 하지만 완전 탐색으로 풀이하면 Time Out으로 풀이에 실패한다. 더 빠르게 문제를 해결할 수 있는 효율적인 코드를 찾아야 한다.

 

정렬후 투 포인터 알고리즘을 사용하면 문제를 해결할 수 있다. 우선 정렬을 한 후, 각 배열의 원소를 축으로 한번씩 둔 다음, 축의 바로 다음 원소를 left 포인터, 배열의 마지막 원소를 right 포인터로 둔 뒤, 합이 0보다 작으면 left 포인터를 전진, 0보다 크면 right 포인터를 후진 시키는 식으로 O(N^2)만에 문제를 해결할 수 있다.

---

소스 코드

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        res = []
        nums.sort()
        
        for i, a in enumerate(nums):
            # same value as before so continue
            if i > 0 and a == nums[i-1]:
                continue
                
            l,r = i+1, len(nums)-1
            while l < r:
                threeSum = a + nums[l] + nums[r]
                if threeSum > 0:
                    r -=1
                elif threeSum < 0:
                    l+= 1
                else:
                    res.append([a,nums[l],nums[r]])
                    l += 1
                    while nums[l] == nums[l-1] and l<r:
                        l+= 1
        return res