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
else:
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
Algorithm
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.
Iterators
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.
Code
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)]
Explore related questions
See similar questions with these tags.