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.

a + b + c = 0

Solution is to just convert this into two pointer technique of solving 2SUM problem. Traverse the array and for each element check if rest of the array has 2 elements that sum up to -a.

class Solution(object):
    def threeSum(self, nums):
        def two_pointer(nums, target):
            '''two pointer technique and it is caller responsibility to pass proper nums'''
            first = 0
            second = len(nums) - 1
            two_sums = []
            while first < second:
                two_sum = nums[first] + nums[second]
                if two_sum < target:
                    first += 1
                elif two_sum > target:
                    second -= 1
                    two_sums.append([nums[first]] + [nums[second]])
                    while first+1 < len(nums) and nums[first] == nums[first+1]:
                        first += 1
                    while second-1 >= 0 and nums[second] == nums[second-1]:
                        second -= 1
                    first += 1
                    second -= 1
            return two_sums
        :type nums: List[int]
        :rtype: List[List[int]]
        nums = sorted(nums)
        three_sums = []
        i = 0
        while i < len(nums)-2:
            two_sums = two_pointer(nums[i+1:], -nums[i])
            for j in range(len(two_sums)):
                three_sums.append([nums[i]] + two_sums[j])
            while i+1 < len(nums) and nums[i] == nums[i+1]:
                i += 1
            i += 1
        return three_sums

1 Answer 1



I find one problem here, you create len(nums)-2 times copy of array by calling nums[i+1:]. It's fast but not necessary. In general helps islice iterator to avoid copies of sub-arrays.


Lines index_of_list_element += constant in loop strongly suggest rewrite code to use proper iterators to get more python way code. In your code i, first, second are indexes for unique numbers in nums. two_pointer() function also can be iterator to avoid create temporary list inside main loop.


Idea of code with iterators

def three_sum(nums):
    def uniques(nums_sorted, start, stop, step):
        prev = None
        for idx in range(start, stop, step):  # xrange() in python2
            value = nums_sorted[idx]
            if value != prev:
                prev = value
                yield idx, value

    def two_pointer(nums_sorted, neg_a, start):
        biter = uniques(nums_sorted, start, len(nums_sorted), 1)
        citer = uniques(nums_sorted, len(nums_sorted) - 1, start, -1)
        (bidx, b), (cidx, c) = next(biter), next(citer)
        while bidx < cidx:
            two_sum = b + c
            if two_sum == neg_a:
                yield b, c
            if two_sum <= neg_a:
                bidx, b = next(biter)
            if two_sum >= neg_a:
                cidx, c = next(citer)

    nums = sorted(nums)
    return [[a, b, c] for aidx, a in uniques(nums, 0, len(nums) - 2, 1)
                      for b, c in two_pointer(nums, -a, aidx + 1)]

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.