파이썬 알고리즘 인터뷰 9번 세 수의 합 (리트코드 15번)

Problem

LeetCode - The World's Leading Online Programming Learning Platform

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이 되는 조합을 찾아라. 중복 X

Solution

1차 풀이 브루트 포스 → time out

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        result = []
        nums = sorted(nums)
        for i in range(len(nums)):
            for j in range(i + 1, len(nums)):
                for k in range(j + 1, len(nums)):
                    if nums[i] + nums[j] + nums[k] == 0 and [nums[i], nums[j], nums[k]] not in result:
                        result.append([nums[i], nums[j], nums[k]])

        return result

일단 뇌를 빼고 풀어보자 ⇒ 어림도 없었다.

2차 풀이 combinations → memory exceed

from typing import List
from itertools import combinations

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        result = []
        combinations_list = list(combinations(nums, 3))
        for group in combinations_list:
            if sum(group) == 0 and sorted(group) not in result:
                result.append(sorted(group))

        return result

itertools의 combination을 사용하여 3개씩 뽑아내면서 해당 group의 합이 0인지 확인하고, sorting한 리스트가 결과에 없으면 더했다.

combinations를 사용하면 숫자 3개의 조합을 쉽게 구할 수 있지만, 전체 조합을 담고 있는 combinations_list의 사이즈가 엄청나게 커진다. n개의 숫자가 들어오면 약 $1/6*n^3$ 개의 원소를 갖게 된다.

3차풀이 - 투포인터

from typing import List

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        result = []
        sorted_nums = sorted(nums)
        for i in range(len(sorted_nums) - 2):
            if i > 0 and sorted_nums[i] == sorted_nums[i-1]:
                continue
            left, right = i + 1, len(sorted_nums) - 1
            while left < right:
                sum_3 = sorted_nums[i] + sorted_nums[left] + sorted_nums[right]

                if sum_3 < 0:
                    left += 1
                elif sum_3 > 0:
                    right -= 1
                else:
                    result.append([sorted_nums[i], sorted_nums[left], sorted_nums[right]])
                    while left < right and sorted_nums[left] == sorted_nums[left + 1]:
                        left += 1
                    while left < right and sorted_nums[right-1] == sorted_nums[right]:
                        right -= 1
                    left += 1
                    right -= 1
        return result

책에 나와있는 투 포인터로 해결하는 방법이다.

주어진 배열을 정렬한 뒤 기준점 뒤에 있는 배열에 투 포인터를 두고 기준점 + 투 포인터의 점들을 더해 합을 확인한다. 만약 합이 0보다 크면 left가 오른쪽으로 한칸 이동하고, 0보다 작다면 right가 왼쪽으로 한칸 이동하는 식이다.

이 문제의 경우 수 조합의 중복을 허용하지 않았으므로 이동하는 점이 이전의 값이랑 같은지 확인한 후 같다면 스킵한다. 이 방법을 통해 중복을 피할 수 있다.

투 포인터를 아직 능숙하게 사용하지 못하는 듯 하다. 투 포인터를 잘 활용하면 시간 복잡도를 한 단계 낮출 수 있는으니 계속 체크해야겠다.

Reference

https://www.pythonpool.com/itertools-combinations/

파이썬 알고리즘 인터뷰 p186 ~ 188