[배열] 새 수의 합

* leetcode 15. 3Sum

 

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

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

# 출력
[
  [-1, 0, 1],
  [-1, -1, 2]
]

---

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

# 3개의 포인터를 이용하여 모두 더해서 0이 되는 값들을 찾는다.
def threeSum(nums: List[int]) -> List[List[int]]:
    results = []
    nums.sort()
    
    # n^3 반복
    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(i + 2, 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]))
                    
    result results

해당 풀이법은 타임아웃으로 풀리지 않습니다. (시간 복잡도 O(n^3))

 

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

# i를 축으로 사용하여 for문을 반복하고 나머지를 투 포인터로 계산하는 풀이법
def threeSum(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]
            
            if sum < 0: # 합 값이 0보다 작으면 왼쪽 포인터를 오른쪽으로 옮겨 값을 증가시켜야한다.
                left += 1
            elif sum > 0: # 반대로 합 값이 0보다 크면 오른쪽 포인터를 왼쪽으로 옮겨 값을 감소시킨다.
                right -= 1
            else: # 합이 0인 경우는 정답이므로 결과를 저장.
                result.append((nums[i], nums[left], nums[right]))
                
                # left, right 포인터가 가리키는 값들의 중복을 스킵하기 위해서 아래 과정을 추가.
                while left < right and nums[left] == nums[left + 1]:
                    left += 1
                while left < right and nums[right] == nums[right - 1]:
                    right -= 1
                
                # 중복값들을 모두 스킵한 후 다름 합 계산을 위해 left, right를 한 칸씩 이동.
                left += 1
                right -= 1

    return results