Python 알고리즘 인터뷰 - 문제9 세수의 합

leetcode 15. 3Sum - https://leetcode.com/problems/3sum/description/

 

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

 

예제

입력

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

 

출력

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

 

풀이

먼저 입력값을 정렬하는 것 부터 시작한다. 리스트에 sort 함수를 사용하면 값이 정렬이 된다.

nums.sort()

정렬된 숫자 리스트

 

정렬된 리스트에서 0 ~ N-1까지 순회하는 인덱스 i를 기준으로 L=i+1, R=len(nums)-1 로 설정 후 L, R 사이를 좁혀가며 nums[i] + nums[L] + nums[R] = 0이 되는 조합을 찾는다. 3숫자의 덧셈을 sum이라고 한다.

 

현재 sum = -1 -1 +4 = 2로 0보다 크다. 이 경우 R의 위치를 왼쪽으로 이동하여 sum <= 0 인 위치까지 이동시킨다. sum = 0이면 해당 조합을 결과 리스트에 추가한다. 

R 위치 이동 및 sum = 0인 경우 결과 리스트에 값 추가

추가가 완료되면 L, R을 각각 이동시킨다. 여기서 R 뿐만 아니라 L을 같이 이동시키는 이유는 R의 값이 변경되는 순간 현재 L의 값은 항상 sum != 0이 되기 때문이다. 예를 들어 R이 이동하여 nums[R] = 1이 되면 nums[i] = -1 이므로 nums[L] = 0이 되어야한다. 이는 현재 ums[L] = -1이 조건을 만족하지 않는다는 의미이므로 L,R 을 동시에 이동시켜도 문제없다.

sum = 0일 경우 L, R 포인터 모두 전진

sum < 0 인 경우 L을 전진시키며 sum > 0인 경우와 동일한 순서로 진행된다. 이를 코드로 나타내면 다음과 같다.

...
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]])
        left += 1
        right -= 1

 

주의해야할 사항은 상황에 따라 중복된 조합이 생길 수 있다는 점이다. 다음과 같은 배열에서 sum = 0인 조합 [-2, 1, 3]을 추가한다. 후 L, R을 이동시킨다.

이 후 L, R을 한칸씩 이동시킬 시 sum = 0을 만족하나, 이전에 추가된 조합과 중복된다.

L, R을 한칸씩 이동시켰을 때 이전 조합과 동일한 값이 추가됨

이를 해결하기 위해 L, R의 다음 값이 동일하면 위치를 계속 이동하도록 한다. 코드로 나타내면 다음과 같다.

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

(좌)중복값 방지를 위한 while loop 내에서의 L, R 이동, (우)다음값 이동을 위한 L,R 이동

 

현재 i에 대한 L, R 조합을 찾는 것을 완료했으면 i를 전진하여 동일하게 진행한다. 여기서 주의할 점은 nums[i - 1] = nums[i]일 경우 중복된 값이 발생한다는 점이다.

i > 0, nums[i] == nums[i-1]인 경우 i 스킵하도록 하며 코드는 다음과 같다.

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

 

전체 코드는 다음과 같다.

 

전체 코드

def threeSum(self, nums):
    results = []
    nums.sort()
    for i in range(len(nums)):
        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:
                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