[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.

 

배열을 입력받아 합으로 0을 만들 수 있는 3개의 엘리먼트를 출력하라.

중복이 있으면 안됨.

 

 

입출력 예시

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

 

생각과정

하나를 기준점으로 잡고 나머지 엘리먼트들 중 두개를 골라 더하여 기준점이 되는 엘리먼트랑 합쳐본다.

그 값이 0이 되면 return하고 아니면 다시 돌아가는 방식을 생각했는데, 접근 방식은 얼추 맞았으나 나는 양 끝의 두 숫자를 한칸씩 안으로 좁혀가면서 더하려고 해서 엉뚱쌩뚱한 답을 리턴했다. 

자, 풀이를 보자.

 

풀이

첫번째 풀이방식은 내가 생각한 방식과 비슷하다. 중첩 for문을 3번 사용하는 방식인데, 문제는 이렇게 풀었을 때 리트코드에서 time limit exeeded를 뱉어낸다. 복잡도가 O(n^3)이기 때문이다. 그래도 코드를 보면서 어떻게 풀이가 되었는지 확인해보자

 

# 첫번째 방식 - 중첩 for문 사용하기 (브루트 포스) -> 리트코드에서 time limit exeeded

nums.sort()

results = []
# i는 index(3)까지 이동함. (뒤에 숫자 2개가 남아있어야하니까)
for i in range(len(nums) - 2):
    # j는 index(4)까지 이동함
    # 중복제거
    if i > 0 and nums[i] == nums[i - 1]:
        continue
    for j in range(i + 1, len(nums) - 1):
        if j > i + 1 and nums[j] == nums[j - 1]:
            continue
        for k in range(j + 1, len(nums)):
            if k > j + 1 and nums[k] == nums[k - 1]:
                continue
            if nums[i] + nums[j] + nums[k] == 0:
                results.append([nums[i], nums[j], nums[k]])

return results

우선 비교를 편하게 하기 위해 정렬을 해준다. 숫자가 막 뒤섞여 있는 것보다는 [-4, -1, -1, 0, 1, 2] 이렇게 정렬되어있는 게 훨씬 값을 비교하기 쉽고, 중복을 찾기도 편하기 때문이다. 

 

정렬을 하지 않고 풀이하는 방식을 찾아보려 했으나 대부분이 정렬하는 방식을 사용했다. 여기서 내가 말하고 싶은 부분은 코테는 수능처럼 누가누가 유형을 잘 파악하고 암기를 잘 하느냐이기 때문에 가장 효율적인 방식을 똑같이 사용한다고 생각한다. 응용편은 알고리즘 중~고수가 됐을 때 알아보고, 지금은 최선의 방식을 공부하자

 

코드 자체가 어렵지 않아서 크게 설명해놓을 부분이 없지만 포인트는 각 for문의 범위가 i는 len(nums) - 2, j는 len(nums) - 1, k는 len(nums)를 사용한다는 점이다.

 

생각해보자 

 

정렬된 [-4, -1, -1, 0, 1, 2] 리스트가 있다.

우리가 찾을 숫자는 총 3개이다. 그럼 첫번째 for문이 가질 수 있는 가장 마지막 인덱스는 3이다. 맨 뒤에 두개의 숫자가 무조건 와야하기 때문이다. 같은 맥락으로 두번째 for문은 뒤에 무조건 한개의 숫자가 와야하고, 마지막 for문은 자신이 마지막 숫자이다. 

 

두번째 고려해야할 사항은 중복인 숫자를 제거하는 부분이다. 문제에 애초에 중복이 오면 안된다고 쓰여있기 때문에 중복을 제거해주는 continue 코드가 들어있다. 참고로 중복은 모든 for문에서 제거해주어야한다.

if i > 0 and nums[i] == nums[i - 1]:
	continue

 

 

이번엔 시간초과가 되지않은 투포인터를 사용한 방식을 살펴보자.

 

nums.sort()

results = []
for i in range(len(nums) - 1):
    # 중복 제거
    if i > 0 and nums[i] == nums[i - 1]:
       continue

    # 투 포인터 설정
    left, right = i + 1, len(nums) - 1

    sum = nums[i] + nums[left] + nums[right]

    # 합계가 0보다 클 때, 제일 큰 수인 nums[right]을 왼쪽(작은 수)로 옮긴다.
    if sum > 0:
        right -= 1
    # 합계가 0보다 작을 때, 제일 작은 수인 nums[left]를 오른쪽(큰 수)로 옮긴다.
    elif sum < 0:
        left += 1
    # 합계가 0일 때, results에 넣는다. 이때,
    else:
        results.append([nums[i], nums[left], nums[right]])

return results

 

여기서 파이썬의 편리한 기능을 볼 수 있다. 바로 '다중할당'이다. 

보통 다른 언어에서 변수를 선언할 때에는 이렇게 일일히 선언해줘야한다.

a = 1;
b = 2;

 

하지만 파이썬에서는 이 부분을 다중할당으로 한줄로 줄일 수 있다.

a, b = 1, 2

따봉 파이썬아 고마워!

 

 

이번 코드에서 사용하는 투포인터는 양쪽 끝에서부터 점차 좁혀나가는 방식이다.

그러므로 다음과 같이 포인터를 설정해준다.

# 투 포인터 설정
left, right = i + 1, len(nums) - 1

 

이제 sum값을 조건문으로 0이 됐는지 비교하면서 포인터를 옮기는 것만 남았다! 어렵지 않은 부분이라 다시 코드를 붙여넣지는 않겠다.

 

조건문이 끝나면 results에 숫자들을 담아주고 return해주면 해결!

---

마치며

투포인터 개념이 점점 익숙해진다! 복습할 수록 발전하는게 눈에 보여서 기분이 좋다.