7장_9 배열(세 수의 합)

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

1. 브루트 포스로 계산

from typing import List


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

        # 브루트 포스 n^3 반복
        for i in range(len(nums) - 2):
            # 중복된 값 건너뛰기
            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

편의를 위해 먼저 정렬을 해서 조금이라도 문제를 간소화 시킨다. 그리고 각각의 3포인터가 이동을 하면서 0인 값을 찾아내는데 중복된 값은 건너뛰어야 하기 때문에 3포인터에 중복된 값은 건너 뛰는 for문을 사용했다.

하지만 이렇게 계산을 해보면 타임아웃으로 인해 실패한다. 시간 복잡도가 O(n세제곱)이기 때문이다. O(n제곱)이내로 계산을 최적화 해보는 방식으로 해 보았다.

2. 투 포인터로 합 계산

from typing import List


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

        for i in range(len(nums) - 2):
            # 중복된 값 건너뛰기
            if i > 0 and nums[i] == nums[i - 1]:
                continue

            # 간격을 좁혀가며 합 `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`인 경우이므로 정답 및 스킵 처리
                    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

        return results

i의 부분을 축으로 중복된 값을 건너 뛰는 방식은 브루트포스로 계산하는 위의 식과 동일하다.
이제 중복이 아닌 경우 투포인터로 계산을 할 수 있다.

앞의 풀이와의 차이점은 시작점이다. i의 다음 지점과 마지막 지점을 left, right로 설정하고 점점 이동하며 합이 0인 값을 찾는 방식이다. sum이 0보다 작다면 세 값의 합을 키워야 하니 작은 쪽인 left를 오른쪽으로 이동시키고 0보다 크다면 세 값의 합을 줄여야 하니 right의 값을 왼쪽으로 이동시키면서 값을 찾는다.

그리고 합이 0인 값을 찾았다면 result에 넘기고 나머지들은 스킵하고 끝이 난다.