m1ndy5's coding blog

LeetCode 15. 3Sum with Python 본문

알고리즘 with python/알고리즘 스터디

LeetCode 15. 3Sum with Python

정민됴 2024. 1. 3. 17:34

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

1차 시도

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        answer = []
        check_multiple = []
        def bt(pair, curr, l):
            if len(pair) == 3:
                s = set(pair)
                if sum(pair) == 0 and s not in check_multiple:
                    answer.append(pair[:])
                    check_multiple.append(s)
                return

            for i in range(curr, l):
                pair.append(nums[i])
                bt(pair, i + 1, l)
                pair.pop()

        bt([], 0, len(nums))

        return answer

백트래킹으로 가능한 조합을 다 찾아보려고 했으나 시간 초과가 났다.
중복이 많을 수 있기 때문에 처음부터 중복을 최대한 피하는 방향으로 가야할 것 같다.

Final Code

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

            # 현재의 오른쪽만 고려하기 때문에 중복 제거 가능
            left, right = i+1, len(nums)-1

            # 경우의 수 다 찾기
            while left < right:
                sum = nums[i] + nums[left] + nums[right]
                # 0보다 작다는건 더 큰걸 더해야한다는 뜻
                if sum < 0:
                    left += 1
                # 0보다 크다는 건 더 작은걸 더해야한다는 뜻
                elif sum > 0:
                    right -= 1
                # 0이 되면
                else:
                    # 결과에 추가
                    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

정렬을 한 뒤 투포인터를 사용한 방법이다.
정렬을 하고 한 번 기준이 된 값은 앞으로 쓰이지 않기때문에 중복없이 찾을 수 있는 방법이었다.