15. 3Sum

[zeikar]

Problem link

https://leetcode.com/problems/3sum/

Problem Summary

배열에서 3 원소의 합이 0이 되는 수를 중복 제외하고 전부 출력하는 문제.

Solution

차례로 i, j, k라 하면 i를 고정할 시 2sum 문제가 된다.

이를 이용해 nums[j] + nums[k] == nums[i] 가 되는 j, k를 구할 수 있고 j 를 쭉 돌면서 -nums[j] 값인 k 가 존재하는지 보면 된다. 이는 2sum 문제처럼 해시맵으로 간단하게 구할 수 있다. 문제는 중복 제거인데 정렬 후 중복된 부분을 제거하면 된다.

위보다 더 빠른 방법이 존재하는데 투 포인터를 사용하는 것이다. 중복 처리도 꽤 간단하고 속도도 빠르다.
nums[i], nums[j], nums[k] 의 합을 보고 합이 음수면 j 를 증가, 양수면 k를 감소시키면서 투 포인터를 적용 가능하다. 중복 제거를 위해 i, j, k 값을 스킵해서 탐색하면 된다.

Source Code

from typing import List


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

        nums.sort()

        for i in range(len(nums)):
            if i > 0 and nums[i] == nums[i - 1]:
                continue

            j = i + 1
            k = len(nums) - 1

            while j < k:
                s = nums[i] + nums[j] + nums[k]

                if s == 0:
                    triplet = [nums[i], nums[j], nums[k]]
                    ans.append(triplet)

                    while j < k and nums[j] == nums[j + 1]:
                        j += 1
                    j += 1

                    while j < k and nums[k] == nums[k - 1]:
                        k -= 1
                    k -= 1

                if s < 0:
                    j += 1

                if s > 0:
                    k -= 1

        return ans