Problem
I’m trying to implement a version of the 3SUM problem. Any help with code style and efficiency would be greatly appreciated. I’m trying to get it roughly O(n2)$O({n}^{2})$. Also, there should be no duplicates. I specifically wanted to try and build upon 2SUM. I know there are other, maybe better, ways to do it.
def two_sum(nums, target):
d = {}
for i in range(len(nums)):
key = target  nums[i]
if key in d:
return [nums[d[key]], nums[i]]
else:
d[nums[i]] = i
return []
def three_sum(nums):
ans = set()
for n in nums:
lst = two_sum(nums, n)
if lst:
lst.append(n)
ans.add(tuple(sorted(lst)))
return ans
print(three_sum(nums))
Solution
1. Bug
The algorithm in three_sum
is not correct (at least not in the usual version of the 3SUM problem):
>>> three_sum([1, 2])
{(2, 1, 1)}
In the call two_sum(nums, n)
there’s nothing to stop two_sum
from using the number n
in its pair, even though n
is also needed to make up the triple.
2. Review

two_sum
andthree_sum
have similar names, so you might guess that they do similar things, buttwo_sum
return at most one pair adding up to an a given target whereasthree_sum
returns a list of triples adding up to zero. When writing a group of similar functions it’s worth taking care that they have similar interfaces — this makes them easier to use reliably.In this case the difference in implementations means that it’s not clear exactly what
three_sum
does — we can’t say that it returns all the triples adding up to zero, because when it callstwo_sum
it gets at most one result, and this might cause it to miss some triples. This is quite unsatisfactory — we can’t use this function if we want all the triples, but if we only want one triple then we have to wait longer than we need to.(This observation is similar in spirit to the zero one infinity rule — we often need functions that find one thing, or find all the things, but very rarely do we need functions that find some of the things.)

There are no docstrings. What do these functions do? What do they return? (I think that if you had written docstrings then you might have spotted the problem in point 1 above.)

In
two_sum
, it’s not clear what the named
means. (I guess d stands for dictionary, but a dictionary containing what?) I would use a more descriptive name likeseen
. Similarly, instead ofkey
I would use a name likeremainder
.