[Algorithm] LeetCode - 3Sum

📠 문제

• 3Sum
• 난이도: Medium
• 배열중 3개를 뽑아 더하면 0이 되는 배열을 겹치지 않게 뽑는 문제

Input: nums = [-1,0,1,2,-1,-4]
Output: [[-1,-1,2],[-1,0,1]]
Explanation: 
nums[0] + nums[1] + nums[2] = (-1) + 0 + 1 = 0.
nums[1] + nums[2] + nums[4] = 0 + 1 + (-1) = 0.
nums[0] + nums[3] + nums[4] = (-1) + 2 + (-1) = 0.
The distinct triplets are [-1,0,1] and [-1,-1,2].
Notice that the order of the output and the order of the triplets does not matter.

🦖 Brute Force

항상 문제를 보면 단순하게 한번 풀어본다. 아래의 풀이는 단순하게 O(n^3)으로 푸는 방법인데 당연하게도 시간초과가 난다. (nums.length는 최대 3,000이라 단순 연산으로 최대 27,000,000,000번 연산을 하게 될 수 있으니까... )

func threeSum(_ nums: [Int]) -> [[Int]] {
    let nums = nums.sorted()
    var triplets: [[Int]] = []

    for i in 0..<nums.count {
        for j in i+1..<nums.count {
            for k in j+1..<nums.count {
                let triplet = [nums[i], nums[j], nums[k]]
                if triplet.reduce(0, +) == 0 && !triplets.contains(triplet) {
                    triplets.append(triplet)
                }
            }
        }
    }

    return triplets
}

💡 풀이

이문제는 좌표 하나를 기반으로 정렬을 한 후 투포인터 알고리즘을 쓰면 된다. 이렇게 된다면 O(n^2)이다. 살짝 까다로운것이 배열이 중복되면 안되기 때문에 똑같은 숫자로 시작하거나 0이 나왔을 때도 똑같은 숫자를 소거해줘야 한다.

func threeSum(_ nums: [Int]) -> [[Int]] {
    let nums = nums.sorted()
    var triplets: [[Int]] = []

    for i in nums.indices {
        /// 똑같은 숫자면 투포인터를 하지 않는다.
        if i > 0 && nums[i] == nums[i-1] { continue }

        var left = i + 1
        var right = nums.count - 1

        while left < right {
            let sum = nums[i] + nums[left] + nums[right]

            if sum == 0 {
                triplets.append([nums[i], nums[left], nums[right]])
               
                /// 0이 나오면 똑같은 숫자를 건너뛴다.
                while left < right && nums[left] == nums[left + 1] { left += 1 }
                while left < right && nums[right] == nums[right - 1] { right -= 1 }

                left += 1
                right -= 1
            } else if sum < 0 {
                left += 1
            } else {
                right -= 1
            }
        }
    }

    return triplets
}

결과는 아래와 같이 준수한(?) 결과가 나왔다.