# Leetcode 3sum problem solution

The problem is:

Given an integer array nums, return all the triplets [nums[i], nums[j], nums[k]] where i, j and k are distinct and nums[i] + nums[j] + nums[k] == 0.

The solution set must not contain duplicate triplets and the order of the output and the order of the triplets does not matter.

This is my code:

def threeSum(nums):
triples = set() # this is where the triples will be stored
num_map = {}
n = len(nums)

for i in range(n): # creating a hash map of the array for quick lookups
num_map[nums[i]] = i

for i in range(n-1):
for j in range(i+1,n):
comp = - nums[i] - nums[j] # the complement
if (comp in num_map and num_map[comp] != i and num_map[comp] != j):
return triples


My questions are:

1. I believe the time complexity of my solution is O(n^2), and it is not known if there is an algorithm which can perform better than O(n^2). Am I correct?
2. In any event, it would seem that my code is not efficient. Which part of the code is slowing it down?

I suspected the use of sorted() might be adding to it, but given that it is only sorting an array of 3 elements (and is only triggered if a solution is found), I thought it might not be having much impact. The reason I sort the array is to avoid duplicates.

• "it would seem that my code is not efficient" - What exactly makes you think that? Mar 3 at 13:49
• @nocomment When I submit it on leetcode leetcode.com/problems/3sum/submissions/1192682909 the runtime is 6553 ms. Most of the submissions run in under 1000ms Mar 3 at 14:03

It would be nice to see some tests of the function. That would help reviewers see what's known to work, and identify untested scenarios.

This transformation is strongly tied to the specific requirements of the question:

for i in range(n): # creating a hash map of the array for quick lookups
num_map[nums[i]] = i


We're storing only a single index in this reverse lookup table, and overwriting any previous value. The wording of the question permits this because we are required to return a set of the values. But this approach makes the code less reusable for similar problems (such as those where we must return the indices rather than values).

for i in range(n) is unidiomatic. I'd write:

    for i,n in enumerate(nums):
num_map[n] = i
print(num_map)


Or more likely:

    num_map = {n:i for i,n in enumerate(nums)}


Since we don't need to know the actual indices used (only that they be distinct), we can transform the input by sorting it input first, thus allowing us to know that nums[i]nums[j]nums[k] when ijk. Then we don't need sorted when adding, and our condition can be simpler:

    nums = sorted(nums)
num_map = dict((n,i) for i,n in enumerate(nums))

for i,nᵢ in enumerate(nums[:-1]):
for j,nⱼ in enumerate(nums[i+1:], i+1):
comp = -nᵢ -nⱼ # the complement
if comp in num_map and j < num_map[comp]:


We can then reduce work by observing that when we reach nⱼ that would make k too small, then no higher j will succeed, so we can break out of the j loop early:

    for i,nᵢ in enumerate(nums[:-1]):
for j,nⱼ in enumerate(nums[i+1:], i+1):
comp = -nᵢ -nⱼ # the complement
if nⱼ > comp:
# we won't find any more matches in this j or higher
break
try:
if j < num_map[comp]:
except KeyError:
pass


I also eliminated the duplication of comp in num_map and num_map[comp] here - remember that Python exceptions are cheap.

We could probably go further, and remove values from comp as they become unreachable. Or maintain an index into it, searching backwards as j increases, rather than doing the in test each iteration. I'll leave that as an exercise for the motivated reader.

• Thank you for the response - this is very helpful. According to wiki.python.org/moin/TimeComplexity the "average case" for a lookup in a set is O(1). Why is it in fact O(log n) in this scenario? Mar 3 at 11:40
• My error - confused with C++ std::set. Mar 3 at 12:38
• I'd say num_map[comp] rather frequently doesn't exist, so raising and catching so many errors seems like a very bad idea, given that they're interested in speed and catching exceptions is expensive. Mar 3 at 14:05

The excellent answer by user @Toby Speight addresses the main goal of performance.

I'm not that smart, but I'll offer some feedback on more mundane style issues with your code.

## Overview

The code layout is nice and compact.

## Documentation

It would be nice to add a docstring to the function to describe ts purpose, what the input is and what it returns. The problem statement in your question would be a good starting point.

The following comment is not very useful because it does not tell the user anything new; it is obvious you are storing something into a variable.

# this is where the triples will be stored


A "hash map" in Python parlance is a "dictionary".

## Naming

The variable triples could be named triplets to be consistent with the documentation.

## Lint check

pylint identified trailing whitespace on a couple lines. Since it is not needed and can be quite annoying, just delete it.

While we're on the topic, I think adding a single space after each comma makes the code easier to read.

Here is new code with the suggestions above:

def threeSum(nums):
'''
Given an integer array nums, return all the triplets
[ nums[i], nums[j], nums[k] ] where i, j and k are distinct
and nums[i] + nums[j] + nums[k] == 0.
<MORE DETAILS TO BE ADDED BY CODE AUTHOR>
'''
triplets = set()
num_map = {}
n = len(nums)

for i in range(n): # creating a dictionary of the array for quick lookups
num_map[nums[i]] = i

for i in range(n-1):
for j in range(i+1, n):
comp = - nums[i] - nums[j] # the complement
if (comp in num_map and num_map[comp] != i and num_map[comp] != j):
return triplets

• Your docstring suggests that threeSum([0,0,0,0]) returns 24 triplets. Mar 3 at 15:27
• thank you for the feedback. I'll take this on board. Mar 3 at 16:31
• @nocomment: Please recommend a better docstring. Go ahead and suggest an improvement by clicking the "Edit" link. My suggestion was really meant as a placeholder for the code author to fill in. I updated the answer to make that more evident. Mar 4 at 11:40

I'm surprised no one on Leetcode mentioned: Assuming you are working with integers within a fixed range [-N,N], CLRS gives a O(n + N log N) solution for 3SUM.

1. Represent nums as a bitset of length N.
2. Here is the trick: compute nums + nums (all pair sums) in O(N log N) time using FFT. If it helps, imagine the bitset as a generating function, and the convolution will add the exponents.
3. Check nums + nums against nums.

3SUM only checks if a solution exists. There are up to n^2 possible triplets to be returned. Otherwise O(n^2) is the best known algorithm. It's simple: for every nums[i], nums[j], simply check if nums[k]=-(nums[i]+nums[j] is in a hash table. If hashing is not allowed, then you can use the hi and lo pointer method like in 2sum to avoid binary search.

• I haven't either :P but I know it exists from reading other people's project euler solutions and cp-algorithms.com
– qwr
Mar 4 at 17:26