[알고리즘] Leetcode #15 세 수의 합 (Python)

 

---

↓↓↓ 아래는 내 리트코드 계정 ↓↓↓

leetcode.com/Jiwon_Lee/

Jiwon Lee - LeetCode Profile

 

---

 

문제 설명

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.

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

• -105 <= nums[i] <= 105

• 0 <= nums.length <= 3000

 

CODE

results = []
nums.sort()

for i in range(len(nums) - 2):
    # continue to duplicate value
    if i > 0 and nums[i] == nums[i - 1]:
        continue

    # calculate sum
    left, right = i + 1, len(nums) - 1
    while left < right:
        sum = nums[i] + nums[left] + nums[right]
        if sum < 0:
            left += 1
        elif sum > 0 :
            right -= 1
        else:
            # sum = 0 (correct case)
            results.append((nums[i], nums[left], nums[right]))

            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

이 문제는 Brute Force로도 풀 수는 있지만, 리트코드 상에서는 O(n^3)으로 타임 아웃된다. 그래서 투 포인터로 접근해야 하는데, 이 코드에서는 i를 고정해두고 left, right를 점점 좁혀가는 방향으로 구현했다. 우선, 입력 값인 배열을 정렬해주고, 중복되는 값에 대해서는 continue 처리를 해준다. 

다음으로, i 다음 값을 left, 맨 마지막 값을 right으로 지정해주고, 세 수의 합이 문제에서 지정한 0이 될 때까지 0보다 작은 경우 left를 늘려 총 합의 크기를 늘려주고, 0보다 큰 경우 right을 줄여 크기를 줄여준다. 마지막으로 합이 0이 되는 순간 results에 해당 값을 할당해주고, 다른 정답이 있는지 탐색한다. 

탐색하기 위해서는 이전에 정답으로 채택되었던 값과 다른 값들 중에 찾아봐야 하므로 while문을 사용해 처리해준 후, left와 right을 동시에 update해준다. (한 쪽만 update해서는 다시 0을 만들 수 없기 때문)

공식 알고리즘은 아니지만 투 포인터 (two-pointer)에 대해 확실히 이해할 수 있게 된 문제!

 

---

 

 

I'm a Senior Student in Data Science ! 

데이터 사이언스를 공부하고 있는 학부생의 TIL 블로그입니다. 게시글이 도움 되셨다면 구독과 좋아요 :)