[Algorithm] LeetCode #15 - 3Sum
개요
LeetCode #15, 3Sum 문제를 풀어봅니다.
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
일단 정렬을 해서 인덱스 i와 i-1 비교를 편하게 할 수 있도록 바꿉니다. 일종의 기준점인 인덱스 i가 인덱스 i-1 값과 같으면 통과합니다. 다루어야 하는 3개 숫자는 인덱스 i, left, right에 있습니다. left는 i+1, right는 len(nums)-1 부터 시작하고, 3개 숫자 합이 0이면 끝나고 합이 0보다 작으면 left를 오른쪽으로 이동시켜 합을 크게 하고 합이 0보다 크면 right를 왼쪽으로 이동시켜 합을 작게 합니다. left와 right에서도 그 다음 순서에서 같은 숫자가 나올 수 있으니, 다른 숫자가 나올 때까지 처리해줍니다.
class Solution:
def threeSum(self, nums: List[int]) -> List[List[int]]:
results = []
# to make it easier to compare i and i +- 1
nums.sort()
for i in range(len(nums)-2):
# skip duplicates
if i > 0 and nums[i] == nums[i-1]:
continue
# move. three numbers at i, left, right
# left starts from i+1, right starts from len(nums)-1
left, right = i + 1, len(nums) - 1
while left < right:
sums = nums[i] + nums[left] + nums[right]
# sorted -> sums < 0 then wanna be large, sums > 0 then wanna be small
if sums < 0:
left += 1
elif sums > 0:
right -= 1
else:
# sum = 0 case
results.append([nums[i], nums[left], nums[right]])
# at least three consecutive equal numbers?
# [-4 2 2 2 2 1 3 5] then -4 2 2 is good, two '2' must be skipped
while left < right and nums[left] == nums[left + 1]:
left += 1
while left < right and nums[right] == nums[right - 1]:
right -= 1
left += 1
right -= 1
return results