3Sum

문제

---

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

3Sum - LeetCode

 

 

나의 풀이

---

리트코드 풀이를 하다가 굉장히 재미있는 문제가 있어 가져와봤다. 문제는 간단하다. 리스트와 타겟이 주어지면 그 리스트의 3개의 원소의 합이 타겟과 같아지는 그룹을 return하는 문제다. 여기서 포인트는 단순히 3중 for문으로 구현하면 시간초과가 난다는 것이다. 다른 방법을 고안해야 하는데 그 풀이가 굉장히 흥미롭다. two-pointer 방식으로 풀어보았다.

1. 주어진 리스트를 정렬한다.
2. i,j,k의 세 포인터를 생성한다. 각각의 역할을 다음과 같다.
   1. i는 시작포인트이다. 범위는 0부터 n-3까지이다. 예를 들어 i가 0이라면 0번째 index의 수는 무조건 합에 포함된다.
   2. j는 i+1부터 시작한다. k는 n-1부터 시작한다.
   3. s = nums[i]+nums[j]+nums[k]의 합을 매번 계산한다. 이 때, 만약 합 s가 target보다 크면 더 작은 수를 더해줘야 하므로 k를 -1 해준다. 반대로, 합 s가 target보다 작으면 더 큰 수를 더해줘야 하므로 j를 +1 해준다.

이 같은 방법을 사용하면 사실상 O(n^3)이 아닌 O(n^2)의 시간복잡도로 계산할 수 있다.(둘 간의 차이는 매우 크다!!)

 

코드

---

class Solution:
            
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        if len(nums) < 3:
            return []
        ans = set()
        nums.sort()
        n = len(nums)
        for i in range(n-2):
            j, k = i+1, n-1
            while j < k:
                s = nums[i] + nums[j] + nums[k]
                if s == 0:
                    ans.add((nums[i], nums[j], nums[k]))
                    j += 1
                    k -= 1
                elif s < 0:
                    j += 1
                elif s > 0:
                    k -= 1
        return ans