부부의 코딩 성장 일기

LeetCode 15(3Sum, Python) 본문

Algorithm/LeetCode

LeetCode 15(3Sum, Python)

제로_콜라 2023. 12. 14. 19:00

1. 문제 링크

 

3Sum - LeetCode

Can you solve this real interview question? 3Sum - Given an integer array nums, return all the triplets [nums[i], nums[j], nums[k]] such that i != j, i != k, and j != k, and nums[i] + nums[j] + nums[k] == 0. Notice that the solution set must not contain du

leetcode.com

2. 문제 설명

  • 주어진 리스트 nums와 서로 다른 인덱스 i, j, k에 대하여 nums[i]+nums[j]+nums[k]=0이 되는 [nums[i], nums[j], nums[k]]를 모두 구하여 한 리스트에 모아 반환하는 문제. 이때 리스트 안의 순서가 다른 것은 고려하지 않음. 
  • 예시) nums = [-1,0,1,2,-1,-4]이 주어지면 [[-1,-1,2],[-1,0,1]]를 반환(-1-1+2=0, -1+0+1=0)

3. 처음 풀이

  • Two Sum 문제 복습(https://codingbubu.tistory.com/1)
    • 이전에 Two sum 문제를 푼 적이 있음. 리스트의 두 수를 더하여 주어진 target num이 되는 두 수를 찾는 문제.(블로그 링크) 이때 이중 for문으로 모든 케이스를 확인하지 않고 딕셔너리에 정보를 저장하며 시간 복잡도 O(n)으로 해결하는 방법이 있었음.
    • 쉽게 말하면 서로 짝꿍인 커플이 있는 지 조사할 때 임의로 두 명씩 뽑아조사하면 오래 걸리지만 앞 사람부터 한 명씩 나와서 앞에 본인의 짝꿍의 이름을 쓰고 들어가면, 본인 차례에 본인 이름이 이미 써져있다면, 그 이름을 쓴 사람 나오라고 하면 훨씬 빨리 짝꿍을 찾을 수 있음. 
    • nums = [2,7,11,15], target = 9인 경우
      • 2에 대해 target의 여분인 9-2=7을 key, 2의 index인 0을 value로 딕셔너리에 저장 seen = {7:0}
      • 그다음 7을 보면 7이 딕셔너리에 있으므로 val = 0번째 2와 더하면 target = 9가 됨. 
      • 이렇게 하면 O(n)으로 Two Sum 해결 가능
    • 이를 이용하여 이중 for문으로 3Sum 해결하기(시간복잡도 O(n²))
      • nums = [-1,0,1,2,-1,-4]
      • -1에 대한 for문만 생각해보자. -1일 때 나머지 두 수를 더하여 1이 되어야하므로 target = 1
      • 0에 대하여 target의 여분인 1-0=1을 key, 0의 index인 1을 value로 딕셔너리에 저장 
        • seen = {1:1}
      • 그 다음 1이 딕셔너리에 있으므로 val = 1이 index인 0과 1을 더하면 target = 1이 되어 -1+0+1=0 가능!
      • 그 다음 2는 딕셔너리에 없으므로 target의 여분인 1-2=-1을 key, 2의 index인 3을 value로 딕셔너리에 저장
        • seen={1:1, -1:3}
      • 그 다음 -1은 딕셔너리에 있으므로 val = 3이 index인 2와 -1을 더하면 target = 1이 되어 -1+2-1=0 가능!
    • 이런 식으로 하면 O(n²) 복잡도로 해결가능하다. 하지만 beat율이 안 좋은 것을 보아 훨씬 빠른 방법이 있나보다. 
class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        nums.sort() #편하게 오름차순 정렬
        seen = {}
        result = []

        
        for i in range(len(nums)): #가장 작은 수 nums[i]에 대해 for문 돌리기
            target = -nums[i] #나머지 두 수를 더하여 -nums[i] 가 되어야함
            seen.clear() #i가 바뀔 때마다 초기화
            for j in range(i+1, len(nums)):
                remainder = target - nums[j] #nums[j]의 짝꿍이 remainder(더해서 target)

                if nums[j] not in seen: #자신 차례에 본인의 이름과 자기 위치를 적고 들어감
                    seen[remainder] = j
                else: #만약 이미 본인 이름이 불렸다면, 짝꿍이 이미 나온 것. 그 짝꿍의 index를 불러내면 됨.
                    ans = [nums[i], remainder, nums[j]]
                    ans.sort() #순서만 다른 답 중복 방지를 위해 정렬 후
                    if ans not in result: #중복되지 않으면 답지에 추가
                        result.append(ans)
        
        return result

 

4. 다른 풀이

  • GPT 풀이
class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        nums.sort() #오름차순 정렬
        result = []

        for i in range(len(nums) - 2): #3개의 수 중 가장 작은 수에 대한 for문
            if i > 0 and nums[i] == nums[i-1]:
                continue  #가장 작은 수가 중복 되는 경우 skip

            left, right = i + 1, len(nums) - 1 #2번 째, 3번 째 수에 대한 two pointers 기법

            while left < right: #가장 작은 수를 고정해두고 나머지 둘을 조정
                curr_sum = nums[i] + nums[left] + nums[right] #현재 합계

                if curr_sum < 0: #음수이면 left 수를 증가
                    left += 1
                elif curr_sum > 0: #양수이면 right 수를 감소
                    right -= 1
                else: #합계 0이면 정답 케이스
                    result.append([nums[i], nums[left], nums[right]]) #답지에 추가

                    while left < right and nums[left] == nums[left + 1]: 
                        left += 1  #left 증가시킬 때 같은 값이면 skip
                    while left < right and nums[right] == nums[right - 1]: 
                        right -= 1  #right 감소시킬 때 같은 값이면 skip

                    left += 1 #합계 0일 때는 (중복제외) left, right 한 쪽만 옮기면 어차피 0이 될 수 없으므로 둘 다 안쪽 이동
                    right -= 1

        return result

  • 다른 사람 풀이
    • 0의 개수에 따라 상황을 나누어 최적화함
    • 시간복잡도는 마찬가지로 O(n²)이지만 더 빠름. 하지만 GPT 풀이가 좀 더 일반적으로 사용 가능.
class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        

        res = set()
        # Split numbers into 3 lists for +ve, -ve and 0s
        n, p, z = [], [], []
        for num in nums:
            if num < 0:
                n.append(num)
            elif num > 0:
                p.append(num)
            else:
                z.append(num)
        
        # Create separate sets for +ve and -ve for fast O(1) lookups
        N, P = set(n), set(p)
        
        # If there are at least 3 zeroes, add 0,0,0
        if len(z) > 2:
            res.add((0,0,0))
            
        # Check is compliment pairs exist. Eg: [-3, 3], then add [-3, 0, 3] provided there is at least 1 zero
        if len(z):
            for num in P:
                if -num in N:
                    res.add((-num, 0, num))
                    
        # Loop through -ve no.s, add 2 -ve no.s and check if the compliment of their sum exists in +ve set. Eg: [-2, -1, 3]
        for i in range(len(n)):
            for j in range(i+1, len(n)):
                target = -(n[i]+n[j])
                if target in P:
                    res.add(tuple(sorted([n[i], n[j], target])))
                    
        #Same as previous step, but reverse case. Eg: [-3, 1, 2]
        for i in range(len(p)):
            for j in range(i+1, len(p)):
                target = -(p[i]+p[j])
                if target in N:
                    res.add(tuple(sorted([p[i], p[j], target])))
        return res

 

5. 배운 점

  • 처음에 two pointers 기법을 해보려다 포기하였음. 그 이유는 하나의 값을 고정시켜두지 않은 채 양쪽 경계를 더하여 0보다 큰 지 작은 지 따져서 그 경계를 옮기는 것을 생각하였는데 이는 빠트리는 경우가 생김.
    • 예시) [-5, 0, 0, 1, 3, 4] 일 때 -5+4<0이므로 l을 오른쪽으로 옮기면 0이 되는데 정답 케이스인 -5, 1, 4를 놓치게 됨. 
  • 하지만 가장 왼쪽 값 하나를 고정하고 나머지 두 개를 찾을 때 two pointers 기법으로 찾으니까 가능. 
  • two pointers등 어떤 방법을 쓸 때 그 방법 하나로만 하려고 하지 말고 복잡도가 조금 올라가더라도 다른 방법과 함께 사용하자.
  • 그리고 시간 복잡도가 같더라도 중복 되거나 필요 없는 경우를 제거하는 최적화에 따라 시간 차이가 꽤 난다. GPT 풀이에서 주석에 skip 설명 있는 부분 생략하면 꽤 느려짐. 시간 복잡도 뿐 아니라 최적화도 꽤 중요함.