LC015 - 3Sum

Problem

Given an integer array nums, return all the triplets [nums[i], nums[j], nums[k]] such that i != j, i != k, and j != k, and nums[i] + nums[j] + nums[k] == 0.

Notice that the solution set must not contain duplicate triplets.

Example

Input: nums = [-1,0,1,2,-1,-4]

Output: [[-1,-1,2],[-1,0,1]]

Explanation: nums[0] + nums[1] + nums[2] = (-1) + 0 + 1 = 0. nums[1] + nums[2] + nums[4] = 0 + 1 + (-1) = 0. nums[0] + nums[3] + nums[4] = (-1) + 2 + (-1) = 0. The distinct triplets are [-1,0,1] and [-1,-1,2]. Notice that the order of the output and the order of the triplets does not matter.

Solution

Naive

Complexity is $O(n^2 \log n)$ time and $O(n)$ space.

def threeSum(self, nums: List[int]) -> List[List[int]]:
	target = 0
	nums.sort()
	s = set()
	output = []
	
	for i in range(len(nums)):
		j = i + 1
		k = len(nums) - 1
		
		while j < k:
			sum = nums[i] + nums[j] + nums[k]

			if sum == target:
				s.add((nums[i], nums[j], nums[k]))
				j += 1
				k -= 1
			
			elif sum < target:
				j += 1
			
			else:
				k -= 1
			
	output = list(s)
	return output

Improved

In theory, this solution would be slower than a two pointer solution, which is in theory$O(n \log n)$. However, in this approach, checking the negative and positive complement while in the worst case is $O(n^2)$, runs faster the majority of times assuming that the ratio of positive and negative values is not skewed because this solution has better constant time factor. Sorting here is effectively $O(1)$ since we sort three elements at a time.

def threeSum(self, nums: List[int]) -> List[List[int]]:

	# if less than three nums, there is no three sum
	if len(nums) < 3:
		return []

	res = set()
	neg = []
	pos = []
	zero = []

	for num in nums:
		if num < 0:
			neg.append(num)
		elif num > 0:
			pos.append(num)
		else: # num == 0:
			zero.append(num)

	posSet = set(pos)
	negSet = set(neg)
		
	# if one zero
	if zero:
		for num in pos:
			if -1 * num in neg:
				res.add((-1 * num, 0, num))
	
	# if three zeroes
	if len(zero) >= 3:
		res.add((0,0,0))
	
	# check positive complement of negative sum
	
	for i in range(len(neg)):
		for j in range(i+1, len(neg)):
			target = -1 * (neg[i] + neg[j])
			if target in posSet:
				# use list so we can sort
				temp = [neg[i], neg[j], target]		
				temp.sort()		
				res.add(tuple(temp))
	
	# check negative complement of positive sum
	
	for i in range(len(pos)):
		for j in range(i+1, len(pos)):
			target = -1 * (pos[i] + pos[j])
			if target in negSet:
				temp = [pos[i], pos[j], target]
				temp.sort()
				res.add(tuple(temp))
	
	return res