LeetCode (4) - 3Sum (파이썬, python)

❓ 문제

• 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: []
  

  Constraints:

  • 0 <= nums.length <= 3000
  • -105 <= nums[i] <= 105

---

Given Code

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        

---

접근법 1

숫자 세개를 더해 합이 0이 되는 케이스를 찾는 문제입니다.

itertools.combinations로 가능한 조합들을 뽑아내 답을 찾으려고 했습니다.

from itertools import combinations
class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        output = []
        candidates = combinations(nums, 3)
        for cand in candidates:
            if sum(cand) == 0:
                cand = sorted(cand)
                if not cand in output:
                    output.append(cand)
        return output

하지만 시간초과가 나서 다른 방법을 찾아봤습니다.

---

접근법 2

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        nums.sort()
        n = len(nums)
        output = []
        for i in range(1, n-1):
            j = 1
            k = 1
            while j + k < n:
                now = nums[i]
                left = nums[i-j]
                right = nums[i+k]
                # print(now, left, right, i, j, k)
                sum = now + left + right
                if sum == 0:
                    local = [now, left, right]
                    if local not in output:
                        output.append(local)
                    j += 1; k += 1
                elif sum < 0:
                    k += 1
                elif sum > 0:
                    j += 1
                if i+k > n-1 or i-j<0:
                    break

        return output

수들을 오름차순으로 정렬한 뒤, now, left, right를 이용해 탐색해가며 조건에 맞는 케이스를 찾도록 프로그래밍하였습니다. 중간에 sum==0일때 j랑 k를 더해주는 부분은 하나의 now에서 left, right에 따라 합이 0인 여러 case가 있는 경우를 위해 추가했습니다.