Find 3 elements that sum to zero in an int array.. I would like to be reviewed on efficiency, style, and obviously if there is a bug I would like to know.
def sum_to_zero(a):
''' given an array of ints, find 3 elements that sum to 0 '''
for i in a:
for j in a[1:]:
for k in a[2:]:
if i + j + k == 0:
return (i,j,k)
There's a bug:
>>> sum_to_zero([4, 17, -2])
(4, -2, -2)
But -2 only appears once in the input, so this is not a legitimate triple.
The docstring is not quite right. "find 3 elements that sum to 0" — and then what? The docstring should say explicitly what the function returns.
The docstring does not say what happens if no three elements sum to zero. In fact, the function returns None. This is a bad idea: it would be all too easy for the caller to forget to check it. Instead, the function should raise an exception.
Instead of three nested loops, use itertools.combinations:
for triple in combinations(a, 3):
if sum(triple) == 0:
return triple
Even when corrected, this function has runtime \$Θ(n^3)\$: it has to examine all possible triples. However, there is a \$Θ(n^2)\$ algorithm:
from collections import defaultdict
from itertools import combinations
class NotFound(Exception):
pass
def sum_to_zero(a):
"""If three elements of the sequence a sum to zero, return them as a
tuple. Otherwise, raise NotFound.
"""
# Map from sum of two elements to a dictionary mapping the index
# of one element to a set of indexes of the other.
pair_sums = defaultdict(lambda:defaultdict(set))
for i, j in combinations(range(len(a)), 2):
d = pair_sums[a[i] + a[j]]
d[i].add(j)
d[j].add(i)
for k, x in enumerate(a):
for i, jset in pair_sums.get(-x, {}).items():
if i != k:
for j in jset:
if j != k:
return a[i], a[j], a[k]
raise NotFound()
