LeetCode：在一个整型数组中找出三个数相加为0的组合

About

今天是挑战LeetCode的第二天，我选择了一道中等难度的题，结果花了3个小时仍然未能通过所有的测试用例，主要是算法时间复杂度太大，超过了时间限制，下面将我的解题过程分享一下

三数之和

描述

给定一个包含 n 个整数的数组 nums，判断 nums 中是否存在三个元素 a，b，c ，使得 a + b + c = 0 ？找出所有满足条件且不重复的三元组。

注意：答案中不可以包含重复的三元组。

例如, 给定数组 nums = [-1, 0, 1, 2, -1, -4]，

满足要求的三元组集合为：

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

解题思路

1. 为了避免答案重复，我先把数组进行了排序
2. 遍历数组，先确定第一个值，求出diff = 0 - firstValue (遍历时如果有重复值就跳过)
3. 然后遍历从第一个选中的值的位置切片后的数组， 求出 diff_another = secondValue - diff (遍历时如果有重复值就跳过)
4. 在剩下的数组中查找 diff_another，如果能找到就说明有解。
5. 为了提高查找效率，我选择了二分查找法。

代码

class Solution(object):
    def binary_chop(self, array, data):
        n = len(array)
        first = 0
        last = n - 1
        while first <= last:
            mid = (last + first) // 2
            if array[mid] > data:
                last = mid - 1
            elif array[mid] < data:
                first = mid + 1
            else:
                return True
        return False

    def threeSum(self, nums):
        """
        :type nums: List[int]
        :rtype: List[List[int]]
        """
        nums.sort()
        ans = []
        for index in range(len(nums)):
            if nums[index] > 0: break
            if nums[index] == nums[index-1] and index > 0: continue
            diff = 0 - nums[index]
            for key in range(index+1, len(nums)):
                if nums[key] > diff: break
                if nums[key] == nums[key-1] and key > index + 1 : continue
                diff_another = diff - nums[key]
                if self.binary_chop(nums[key+1:], diff_another):
                    ans.append([nums[index], nums[key], diff_another])
        return ans

运行结果

通过了311个测试用例，但是当测试用例的元素超过3000后，该算法运行时间超过14s，未能达到要求

参考算法

采用三个指针分别指向左边，中间，右边三个值，在运算的时候选择从两边向中间逼近，大幅减小了时间复杂度。

class Solution:
    def threeSum(self, nums):
        nums.sort()
        res, k = [], 0
        for k in range(len(nums) - 2):
            if nums[k] > 0: break # because of j > i > k.
            if k > 0 and nums[k] == nums[k - 1]: continue # skip.
            i, j = k + 1, len(nums) - 1
            while i < j:
                s = nums[k] + nums[i] + nums[j]
                if s < 0:
                    i += 1
                    while i < j and nums[i] == nums[i - 1]: i += 1
                elif s > 0:
                    j -= 1
                    while i < j and nums[j] == nums[j + 1]: j -= 1
                else:
                    res.append([nums[k], nums[i], nums[j]])
                    i += 1
                    j -= 1
                    while i < j and nums[i] == nums[i - 1]: i += 1
                    while i < j and nums[j] == nums[j + 1]: j -= 1
        return res

参考链接

转载自：https://leetcode-cn.com/problems/3sum/