[파이썬 알고리즘] 세 수의 합

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

ex)

Input: nums = [-1,0,1,2,-1,-4]
Output: [[-1,-1,2],[-1,0,1]]

풀이 1) 브루트 포스로 계산

-> 일일이 하나씩 다 계산해 보는 방법이며 가장 쉬운 방법이다. 일일이 계산하니까 당연히 시간복잡도가 증가한다.

->  앞뒤로 같은 값이 있을 경우, 쉽게 처리하기 위해 sort() 함수를 사용! nums sort() => [-4, -1, -1, 0, 1, 2]

-4	-1	-1	0	1	2
i	j	k	→
i		j	k	→
i			j	k	→

-> 이런 식으로 이동하여 i + j + k =0 을 찾는다.

-> 브루트 포스는 일일이 다 계산하다보니 중복도 있을 수 있어서 continue로 건너뛰기 처리!

if i>0 and nums[i] == nums[i -1]:
    continue

전체코드

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]
print(threeSum(nums))  # [[-1, -1, 2], [-1, 0, 1]]

-> 타임아웃 실패!

 

풀이 2) 투 포인터로 합 계산

-> i를 축으로 하고, 중복된 값을 건너뛰게 한 부분은 앞서 풀이와 동일

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

-> 중복된 값 건너뛰고

-> i를 기준으로 left , right를 정해서 점점 좁히는 방식

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:
                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