[리트코드 15] 투 포인터(Two Pointers) 를 활용한 문제풀이3

*투 포인터에 대해 모른다면, [리트코드 334] 투 포인터(Two Pointers) 를 활용한 문제풀이1 참고.

 

투 포인터를 활용하면 브루트 포스로 풀이하였을 때보다 시간복잡도(O)를 줄일 수 있다.

이번 문제도 브루트 포스로 풀이가 가능하나 시간복잡도O(n^3) 로 타임 아웃이 발생할 것이다. 그럼으로 투 포인터를 활용해 시간복잡도를 줄일 필요가 있다.

 

15. 세 수의 합

Example

- 주어진 숫자 배열(nums) 중 합하여 0이 되는 세 개의 수를 구하라. (단, 정답에 동일한 세트가 있으면 안됨)

예)

[[-1, -1, 2], [-1, 0, 1]] (o)

[[-1, -1, 2], [2, -1, -1]] (x)

Constraints

- 0 <= nums.length  <= 3000

- -10^5 <= nums[i] <= 10^5

Solutions

1. 투 포인터 풀이를 위해 배열을 정렬한다.

2.  nums[i] + nums[j] + nums[k] = 0 일때 i 를 축으로 i+1 을 j 로 끝 지점을 k 로 설정하고 0 ~ (배열의 길이 - 2) 까지 for문을 돌며 세수의 합을 확인한다. (j, k 두 지점을 뺀 범위) 

* 동일한 세트가 정답에 들어있으면 안되므로 i의 값이 이전 값이랑 같은 경우 스킵한다.

3. 이제 두 포인터(j, k) 의 간격을 좁혀가며 세수의 합이 0 보다 작으면 왼쪽 포인터(j)를 오른쪽으로 +1, 0보다 크면 오른쪽 포인터(k)를 왼쪽으로 -1 이동하면서 합이 0이 되는 값을 찾는다.

4. 세수의 합이 0인 값을 찾았으면 정답 리스트에 추가한 후 나머지 정답을 확인하기 위해 두 포인터 모두 움직여 준다. 

* 동일한 세트가 정답에 들어있으면 안되므로 같은 값이 아닐때까지 이동한다. (while 문)

코드 샘플

# python
class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        answer = []
        
        # 정렬
        nums.sort()
        
        # i를 축으로 loop  
        for i in range(len(nums) - 2):
            # 이전 값이랑 같은 경우 스킵한다.
            if i > 0 and nums[i] == nums[i - 1]:
                continue
            
            j, k = i + 1, len(nums) - 1
            
            # 투 포인터로 간격을 좁혀나가며 세수의 합 확인
            while j < k:
                sum = nums[i] + nums[j] + nums[k]
                
                # 세수의 합이 0보다 작으면 왼쪽 포인터를 오른쪽으로 +1
                if sum < 0:
                    j += 1
                    
                elif sum > 0:
                    k -= 1
                    
                else:
                    # 세수의 합이 0인 경우
                    answer.append([nums[i], nums[j], nums[k]])
                    
                    while j < k and nums[j] == nums[j+1]:
                        j += 1
                    
                    while j < k and nums[k-1] == nums[k]:
                        k -= 1
                    
                    # 나머지 정답도 찾기 위해 포인터를 움직인다.
                    j += 1
                    k -= 1
                    
        return answer

 

(번외) 파이썬 set 자료 구조를 이용해 중복 값 제거

코드는 간결해졌지만, 실행 속도는 위의 코드의 3배 가량 느렸다...

# python
class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        answer = set()
        
        nums.sort()
        
        for i in range(len(nums) - 2):
            j, k = i + 1, len(nums) - 1
            
            while j < k:
                sum = nums[i] + nums[j] + nums[k]
                
                if sum < 0:
                    j += 1
                    
                elif sum > 0:
                    k -= 1
                    
                else:
                    answer.add((nums[i], nums[j], nums[k]))
                    j += 1
                    k -= 1
                    
        return [list(a) for a in answer]

 

참조: 파이썬 알고리즘 인터뷰 https://github.com/onlybooks/algorithm-interview