[파이썬 알고리즘 인터뷰] 배열 - 세 수의 합

🐍 문제

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

  • 입력
    nums = [-1, 0, 1, 2, -1, -4]
  • 출력
    [[-1, 0, 1], [-1, -1, 2]]

🐍 내 풀이

이런 방법밖에 생각이 안난다.. 슬프다...

def sum_zero(nums: list) -> list:
    if len(nums) < 3:
        return []

    result = []

    for idx1 in range(0, len(nums) - 1):
        for idx2 in range(idx1 + 1, len(nums) - 1):
            for idx3 in range(idx2 + 1, len(nums) - 1):
                if nums[idx1] + nums[idx2] + nums[idx3] == 0:
                    if sorted([nums[idx1], nums[idx2], nums[idx3]]) in result:
                        continue
                    result.append(sorted([nums[idx1], nums[idx2], nums[idx3]]))

    return result                    


nums = [-1, 0, 1, 2, -1, -4]
print(sum_zero(nums))  # [[-1, 0, 1], [-1, -1, 2]]

🐍 교재 풀이1 - 브루트 포스로 계산

똑같이 모든 리스트의 값을 세번 반복하여 값을 확인하지만 조금이라도 시간을 줄이기 위해 정렬하고 앞뒤로 동일한 값을 제외한다(값이 동일하다면 똑같은 리스트가 나올테니까). 시간복잡도는 n3이며 타임아웃이 발생한다.

def threeSum(nums: list) -> list:
    results = []
    nums.sort()

    # 브루트 포스 n 세제곱 반복
    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

nums = [-1, 0, 1, 2, -1, -4]
threeSum(nums)  # [[-1, -1, 2], [-1, 0, 1]]

🐍 교재 풀이2 - 투 포인터로 합 계산


한번에 왼쪽 값과 오른쪽 값을 이동하는 투 포인터 방식에 i를 추가했다. i위치를 기준으로 왼쪽 값, 오른쪽 값을 받고 세 수의 합을 확인한다.

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

nums = [-1, 0, 1, 2, -1, -4]
threeSum(nums)  # [[-1, -1, 2], [-1, 0, 1]]

세 수의 합이 특정한 값이 되려면 무조건 주어진 리스트를 모두 훑어서 값의 합을 확인해야 한다. 세 개의 수의 합이므로 세 번을 반복해야 할듯한데 세번 반복해야 할 것을 어떻게 줄이느냐가 핵심인거같다. 이번 문제처럼 여러번 반복해야 할것처럼 보이는 문제에서 투 포인터 방식을 사용하면 최소한 한 번의 반복은 줄일 수 있을 것이다.

좋은 웹페이지 즐겨찾기