15. 3Sum

https://leetcode.com/problems/3sum/description/

Description

Given an array nums of n integers, are there elements a, b, c in nums such that a + b + c = 0? Find all unique triplets in the array which gives the sum of zero.

Note:

The solution set must not contain duplicate triplets.

Example:

Given array nums = [-1, 0, 1, 2, -1, -4],

A solution set is:

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

Solution

class Solution(object):
    def threeSum(self, nums):
        """
        :type nums: List[int]
        :rtype: List[List[int]]
        """
        import bisect
        if len(nums)<3: return []
        ret = set()
        nums.sort()
        # 第一次优化, 由于选中的最多是3个元素，故重复元素仅保留 3 个
        tp_nums = nums[:3]
        for i in range(3, len(nums)):
            if not (nums[i]==nums[i-1]==nums[i-2]==nums[i-3]):
                tp_nums.append(nums[i])
        nums = tp_nums
        for i in range(len(nums)):
            for j in range(i+1, len(nums)):
                # 使用二分查找加快查找速度
                k = bisect.bisect_left(nums, -(nums[i]+nums[j]), j+1)
                if k<len(nums) and nums[k]+nums[i]+nums[j]==0:
                    ret.add((nums[i], nums[j], nums[k]))
                elif k==j+1:
                    #  第二次优化
                    break
        return [list(i) for i in ret]

理论上的时间复杂度为 $O(n^2\log(n))$ ，实际上由于各种优化，一定会比这个复杂度要小；