3Sum

问题描述

Given an array S of n integers, are there elements a, b, c in S 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.

For example, given array S = [-1, 0, 1, 2, -1, -4],

A solution set is:
[
  [-1, 0, 1],
  [-1, -1, 2]
]

方案

def three_sum(nums):
    if not nums:
        return []
    length = len(nums)
    triples = []
    nums.sort()
    for i in range(length - 1):
        if i == 0 or (i != 0 and nums[i] != nums[i - 1]):
            s = 0 - nums[i]
            low = i + 1
            high = length - 1
            while low < high:
                tmp = nums[low] + nums[high]
                if tmp == s:
                    triples.append([nums[i], nums[low], nums[high]])
                    while low < high and nums[low] == nums[low + 1]:
                        low += 1
                    while low < high and nums[high] == nums[high - 1]:
                        high -= 1
                    low += 1
                    high -= 1
                elif tmp < s:
                    while low < high and nums[low] == nums[low + 1]:
                        low += 1
                    low += 1
                else:
                    while low < high and nums[high] == nums[high - 1]:
                        high -= 1
                    high -= 1
    return triples

print three_sum([-1, 0, 1, 2, -1, -4])