파이썬 알고리즘 인터뷰 - 세 수의 합

문제

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

 

입력

nums = [-1, 0, 1, 2, -1, -4]

출력

[
    [-1, -1, 2],
    [-1, 0, 1]
]

작성한 코드

def threeSum(nums: List[int]):
    lists = []
    for a in itertools.combinations(sorted(nums), 3):
        if sum(a) == 0:
            lists.append([*a])

    rtn = []
    for l in lists:
        if l not in rtn:
            rtn.append(l)
    return rtn

리코드에서 시간 초과가 나왔다. - 숫자가 많아지면 조합해야 할 숫자가 많기 때문에 초과된걸로 생각된다.

 

책에 있는 코드

def threeSum(nums: list) -> list:
    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

 

리코드에서 확인

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

3Sum - LeetCode Discuss

 

회고

책에 있는 코드가 너무 이해가 안되면 리코드에 디스커션을 확인해보면 다른 코드를 볼 수 있다.

30~40분만 고민해보고 풀리지 않는 다면 답을 보고 이해하려고 노력해야겠다.

책에있는 답도 이해가 되지 않는다면 내가 이해 할 수 있는 답을 구글링해서 찾아보는것도 좋은 방법이라고 생각한다.