3Sum (Two Pointers)

LeetCode의 3Sum 문제다.

3Sum - LeetCode

주어진 수 중에서 3개를 골라 합이 0인 조합을 반환하는 문제인데 중복된 조합은 반환하지 않아야 한다. 정수 배열은 10000개 까지 주어질 수 있기 때문에 최적화되지 않은 코드의 경우 시간 초과가 발생할 수 있다. 실제로 첫 번째로 구현한 아래 코드는 28번째 테스트 케이스에서 시간 초과 에러가 발생했다.

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        self.solutions = []
        
        if len(nums) < 3:
            return []
        
        for index in range(len(nums)):
            self.findSum(nums[:index] + nums[index+1:], [nums[index]])
        
        return self.solutions
    
        
    def findSum(self, nums, triplet):
        if len(triplet) == 3 and sum(triplet) == 0:
            for solution in self.solutions:
                if sorted(solution) == sorted(triplet):
                    return
            
            self.solutions.append(triplet)
            return
        
        for index in range(len(nums)):
            self.findSum(nums[:index] + nums[index+1:], triplet+[nums[index]])

이전에 DFS, BFS를 학습하면서 얻은 지식으로 BFS처럼 시도해보았다. 정수를 3개 골라서 합하여 0을 만들어야 하기 때문에 처음 배열에서 모든 정수를 각각 첫 번째 정수로 삼고 나머지 두 번째, 세 번째 정수를 고르기 위해 첫 번째 정수를 제외한 나머지와 첫 번째 정수가 담긴 리스트를 매개변수로 하여 findSum 메서드를 호출한다.

 

메서드에서는 정수를 고를 배열과 고른 정수를 담을 리스트를 매개변수로 받는다. 그래서 처음 호출됐을 때, 즉 두 번째 정수를 탐색할 때도 비슷하게 배열의 모든 정수를 각각 목록(triplet)에 추가하고 해당 정수를 제외한 나머지 배열을 재귀 함수로 넘겨서 세 번째 정수를 탐색한다. 세 번째 정수도 탐색하고 세 정수를 합한 값이 0이며 기존에 탐색한 배열에 포함되어 있지 않다면 결과 리스트(solutions)에 추가해서 threeSum 메서드에서 이를 반환하는 로직이다.

하지만 얼마 못가 시간 초과가 발생한 것을 볼 수 있다. 정수가 11개밖에 없는데 왜 시간 초과가 발생한 걸까? 이는 최적화가 안됐기 때문이다.

예를 들어 위처럼 [-3,-1,0,1,1,0,-1,3] 같은 정수가 주어졌을 때 눈으로 보기에도 쉽게 찾을 수 있는 조합은 [-1,0,1]이다. 물론 이 조합은 다른 탐색에서도 발견될 수 있기 때문에 중복 검사가 필요하다.

예를 들어 [-1,0,1]의 중복 조합으로는 위처럼 [-1,0,1], [1,0,-1], [1,0,-1]가 있다. 그런데 중복된 조합을 판단한다고 해도 애초에 이 중복된 조합까지 탐색하는 게 문제 아닐까? 게다가 중복된 조합까지 탐색하는 시간과 중복 여부를 탐색하는 시간도 합하면 더 시간이 오래 걸릴 것이다. 그래서 책에서는 일단 배열을 정렬하라고 조언하고 있다.

파이썬 리스트의 정렬 메서드는 O(n*logn)이라는 준수한 성능을 갖고 있으며 문제에서는 주어진 정수 배열의 인덱스를 기억할 필요가 없기 때문에 정렬해도 별 문제가 없다. 게다가 중복을 제거하라는 조건도 붙어있기 때문에 이런 경우 배열을 정렬해서 푸는 게 더 최적화하기 좋은데 왜냐면 위의 그림에서 볼 수 있듯이 같은 원소, 즉 중복된 원소를 쉽게 파악할 수 있기 때문이다. 이걸로 어떻게 중복 탐색 자체를 피할 수 있을까? 원리는 간단하다. 첫 번째 정수를 등록할 때 현재 정수와 이전 정수가 동일하다면 탐색하지 않는 것이다.

만약 정렬된 정수 배열에서 0이 되는 조합을 탐색할 때 위의 예시처럼 인덱스 1에서 시작하는 [-1,0,1]과 인덱스 2에서 시작하는 [-1,0,1]이 있을 때 이 두 조합은 중복된 조합이기 때문에 인덱스 2에서는 이전 정수와 현재 정수가 동일하기 때문에 탐색하지 않고 넘어갈 수 있다. 그럼 만약 중복된 정수가 두 개 필요한 경우는 어떻게 될까?

언급했듯이 현재 정수와 "이전" 정수가 동일하면 탐색을 생략하기 때문에 중복된 정수를 포함하는 조합은 최소한 한 번은 탐색된다. 예를 들어 [-3,-1,-1,-1,1,2] 정수 배열이 있을 때 인덱스 1의 -1을 첫 번째 정수로 탐색한다면 [-1,-1,2] 같은 조합을 구성할 수 있다. 그러나 인덱스 2의 -1은 이전 정수(인덱스 1의 -1)와 동일하기 때문에 굳이 [-1,-1,2] 같은 조합을 탐색하지 않는다. 그러므로 문제없이 중복을 건너뛰면서 탐색할 수 있다.

 

그러나 인덱스 1의 -1을 첫 번째 정수로 선택했을 때 0을 만들려면 인덱스 2와 인덱스 3의 -1 둘 다와 조합될 수 있기 때문에 여전히 중복이 발생한다. 이는 어떻게 해야 할까?

중복된 조합을 탐색하는 것을 해결하기 위해 책에서는 내가 시도했던 BFS 대신 두 개의 포인터를 이용한 풀이를 제시했다. 먼저 첫 번째 정수 이후의 배열에서 맨 처음과 맨 끝을 각각 두 번째, 세 번째 정수라고 두고 탐색한다. 그리고 세 정수를 모두 더한 값이 0보다 작다면 왼쪽의 포인터, 즉 두 번째 정수를 가리키는 포인터를 오른쪽으로 이동시키고 0보다 크다면 오른쪽 포인터, 즉 세 번째 정수를 가리키는 포인터를 왼쪽으로 이동시킨다. 왜 이렇게 이동하는 것일까? 왜냐면 지금 배열은 정렬된 상태기 때문이다.

 

정렬된 정수 배열에서는 왼쪽으로 갈수록 더 작은 수가, 오른쪽으로 갈수록 더 큰 수가 존재한다. 그러므로 만약 합이 0보다 크다면 큰 정수를 가리키고 있는 오른쪽 포인터를 왼쪽으로 이동시키고 0보다 작다면 작은 정수를 가리키고 있는 왼쪽 포인터를 오른쪽으로 이동시켜서 최대한 0에 가깝게 만드는 것이다.

위의 예시에서는 [-1,-1,3]의 합이 1로 0보다 크기 때문에 맨 마지막 3을 가리키는 오른쪽 포인터가 왼쪽으로 한 칸 이동하여 [-1,-1,1] 조합이 된다. 이번에는 합이 -1로 0보다 작기 때문에 인덱스 1 이후의 배열에서 맨 처음 정수를 가리키는 왼쪽 포인터가 오른쪽으로 한 칸 이동하여 [-1,0,1]이 된다. 그래서 0이 되는 조합을 발견할 수 있는 것이다. 그런데 0이 되는 조합을 발견하면 바로 탐색을 종료해야 할까? 그렇지 않다.

위의 예시를 보자. 첫 번째 정수인 -3을 기준으로 0을 만들 수 있는 조합은 [-3,0,3]과 [-3,1,2] 두 가지가 있다. 하지만 [-3,0,3]을 발견하고 바로 탐색을 종료하면 이 [-3,1,2] 조합을 탐색할 수 없다. 그러므로 어떤 조합을 발견했다면 왼쪽 포인터와 오른쪽 포인터를 최대한 안쪽으로 붙여서 다시 탐색해야 한다. 안쪽으로 붙인다는 것은 현재 왼쪽 포인터와 오른쪽 포인터가 가리키고 있는 값이 중복되는 동안 각각 오른쪽과 왼쪽으로 이동해야 한다는 것이다.

중복된다는 것은 위의 그림처럼 현재 포인터가 가리키는 값과 같은 정수를 말한다. 조합 [-3,0,3]을 발견했을 때 왼쪽 포인터는 인덱스 1의 0을 가리키고 있다. 그러나 인덱스 1의 0이든 인덱스 2의 0이든 결과적으로 동일한 조합이기 때문에 생략하고 오른쪽으로 이동하다가 인덱스 4의 정수 1은 새로운 조합이 있을 가능성이 있기 때문에 더 이동하지 않는 것이다. 비슷하게 오른쪽 포인터도 같은 3이 나올 동안 왼쪽으로 이동하다가 인덱스 5의 정수 2에서 멈춘다. 그리고 여기서부터 다시 탐색을 시작해서 0이 되는 조합을 찾아나가는 것이다.

 

이는 두 번째 정수와 세 번째 정수만 활용하는 방안이기 때문에 배열의 모든 요소, 정확히는 두 번째, 세 번째 정수가 첫 번째 정수 뒤에 있으므로 맨 뒤에서 세 번째 정수까지 각 정수들을 첫 번째 정수로 탐색해야 한다. 이를 기반으로 구현한 코드는 다음과 같다.

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        nums.sort()
        results = []
        
        for index, value in enumerate(nums):
            # 맨 뒤에서 3번째 요소까지만 탐색한다.
            if index > len(nums) - 2:
                break
            
            # 현재 요소가 이전 정수와 동일할 경우 탐색하지 않는다.
            if index > 0 and value == nums[index-1]:
                continue
            
            # 왼쪽 포인터, 오른쪽 포인터를 설정한다.
            index_left, index_right = index+1, len(nums)-1
            
            # 두 포인터가 서로를 건드리지 않는 동안 탐색한다.
            while index_left < index_right:
                current_sum = value + nums[index_left] + nums[index_right]
                # 세 정수의 합에 따라 포인터를 이동시킨다.
                if current_sum > 0:
                    index_right -= 1
                elif current_sum < 0:
                    index_left += 1
                else:
                    results.append([nums[index_left], value, nums[index_right]])
                    # 현재 첫번째 정수를 기준으로 추가적인 조합을 탐색한다.
                    while index_left < len(nums)-1 and nums[index_left] == nums[index_left+1]:
                        index_left += 1
                    while index_right < len(nums)-1 and nums[index_right] == nums[index_right+1]:
                        index_right -= 1
                    
                    # 각 포인터가 이동한 후 새로운 정수로 넘어가기 위해 한번 더 이동한다.
                    index_left += 1
                    index_right -= 1
                    
                    
        return results

맨 마지막에 왼쪽, 오른쪽 포인터를 한번 더 이동시키는 것은 while 문에서 현재 정수와 그다음 정수를 비교하여 동일한 정수일 동안만 탐색했기 때문이다. 그래서 동일한 정수의 맨 마지막 정수를 가리키고 반복이 끝나기 때문에 거기서 한번 더 넘어가 줘야 다른 정수를 가리킬 수 있는 것이다.

 

 

쉽지 않은 문제였고 다른 사람의 풀이를 보면 대부분 Counter를 활용한 것 같은데 잘 이해가 가진 않는다. 직관적으로 이해가 가는 건 아무래도 이 두 개의 포인터를 이용한 풀이인 것 같다.