[Q. 09] 세 수의 합

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

 

이 또한 모든 경우의 수를 계산하면 n의 3승으로 매우 오래 걸린다. 그래서 이 문제도 투포인터로 해결한다.

먼저 배열을 정렬한다.

nums.sort()

 

그 다음, 각 원소마다 for문을 돌며 3수의 합을 찾는다.

for i in range(len(nums)-2):
    if i > 0 and nums[i] == nums[i-1]:
        continue

 

i 가 i-1과 같다면 같은 결과가 나올 것이기에 이는 넘어가주는 조건문을 작성한다.

left, right = i+1, len(nums)-1

while left < right: 
    sum = nums[i] + nums[left] + nums[right]

    if sum == 0:
        result = [nums[i], nums[left], nums[right]]
        if result not in answer:
            answer.append(result)
        left += 1
    elif sum < 0:
        left += 1
    else:
        right -= 1

 

그리고 투포인터 설정을 한다. left는 i의 바로 오른쪽 값, right는 배열의 마지막 값으로 설정 한 후 left<right가 될 때 까지 아래 조건을 반복한다.

현재 for문의 i가 가리키는 값 nums[i]와 left/right 포인터가 가리키는 nums 값들의 합이 우리가 구할 세 수의 합이다. 이 세 수의 합이 0이면 우리가 찾는 조건이기에 answer 배열에 추가하는데, 중복된 값이 들어갈 수도 있으므로 not in을 통해 걸러준다.

그리고 sum 값이 0보다 작다면 더 큰 값이 필요하기에 left를 증가시켜주고, sum 값이 0보다 크다면 더 작은 값이 필요하기 때문에 right를 감소시켜준다. 이는 배열이 순차적으로 정렬되어있기 때문에 가능하다.

 

전체 코드는 다음과 같다.

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        nums.sort()
        answer = []
        
        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:
                    result = [nums[i], nums[left], nums[right]]
                    if result not in answer:
                        answer.append(result)
                    left += 1
                elif sum < 0:
                    left += 1
                else:
                    right -= 1

        return answer