Given an array find any three numbers which sum to zero

Question

Looking for optimizations and cleaner, more pythonic ways of implementing the following code.

#Given an array find any three numbers which sum to zero.

import unittest


def sum_to_zero(a):
    a.sort(reverse=True)
    len_a = len(a)
    if  len_a < 3:
        return []
    for i in xrange(0, len_a):
        j = i + 1
        for k in xrange(j + 1, len_a):
            if a[j] + a[k] == -1 * a[i]:
                return [a[i], a[j], a[k]]
    return []


class SumToZeroTest(unittest.TestCase):

    def test_sum_to_zero_basic(self):
        a = [1, 2, 3, -1, -2, -3, -5, 1, 10, 100, -200]
        res = sum_to_zero(a)
        self.assertEqual([3, 2, -5], res, str(res))

    def test_sum_to_zero_no_res(self):
        a = [1, 1, 1, 1]
        res = sum_to_zero(a)
        self.assertEqual(res, [], str(res))

    def test_small_array(self):
        a = [1, 2, -3]
        res = sum_to_zero(a)
        self.assertEqual(res, [2, 1, -3], str(res))

    def test_invalid_array(self):
        a = [1, 2]
        res = sum_to_zero(a)
        self.assertEqual(res, [], str(res))

#nosetests sum_to_zero.py

Answer 1

Your code fails this test case:

def test_winston1(self):
    res = sum_to_zero([125, 124, -100, -25])
    self.assertEqual(res, [125,-100,-25])

As for the code

def sum_to_zero(a):
    a.sort(reverse=True)

Why sort?

    len_a = len(a)
    if  len_a < 3:
        return []
    for i in xrange(0, len_a):

You can just do xrange(len_a). I recommend doing for i, a_i in enumerate(a) to avoid having the index the array later.

        j = i + 1

Thar be your bug, you can't assumed that j will be i + 1.

        for k in xrange(j + 1, len_a):
            if a[j] + a[k] == -1 * a[i]:

Why this instead of a[j] + a[k] + a[i] == 0?

                return [a[i], a[j], a[k]]
    return []

Not a good a choice of return value. I'd return None. An empty list isn't a natural continuation of the other outputs.

Here is my implementation of it:

def sum_to_zero(a):
    a.sort(reverse = True)
    for triplet in itertools.combinations(a, 3):
        if sum(triplet) == 0:
            return list(triplet)
    return []

Answer 2

1. Comments on your code

In addition to the improvements Winston made in his answer, I would add:

1. Your function needs a docstring. What does it do and how do you call it?

2. It's better to raise an exception rather than return an exceptional value as a result. It's too easy for the caller to omit the check for the exceptional value.

3. The number zero is arbitrary, so why not generalize the code so that can find triples that sum to any target number?

4. It's not clear to me what the purpose of a.sort(reverse=True) is. This has the side-effect of sorting the input a into reverse order, which is probably not what the caller expected. If you really need to process a in reverse order, you should write something like:

   a = sorted(a, reverse=True)
   

   so that the input to the function is unchanged. But here there doesn't seem to be any reason for the sort, so you could just omit it.

So my implementation of the test-all-triples algorithm looks like this:

from itertools import combinations

class NotFoundError(Exception):
    pass

def triple_with_sum(a, target=0):
    """Return a tuple of three numbers in the sequence a that sum to
    target (default: 0). Raise NotFoundError if no triple sums to
    target.

    """ 
    for triple in combinations(a, 3):
        if sum(triple) == target:
            return triple
    raise NotFoundError('No triple sums to {}.'.format(target))

2. A better algorithm

The test-all-triples algorithm is very straightforward, and runs in O(n³) time and O(1) extra space.

However, it's worth noting that it is possible to solve the problem in O(n²) time and O(n) extra space, using the following algorithm:

1. Store, for each number in the input, the set of positions in the input at which that number can be found. This step takes O(n) time and uses O(n) extra space.

2. For each pair of distinct positions i, j in the input, consider the negation of the sum of the values of those two positions. If this appears at some position in the input (other than position i or j) then we have three numbers in the input that sum to 0. This step takes O(n²) time.

Here's a straightforward implementation:

from collections import defaultdict

def triple_with_sum2(a, target=0):
    """Return a tuple of three numbers in the sequence a that sum to
    target (default: 0). Raise NotFoundError if no triple sums to
    target.

    """ 
    positions = defaultdict(set)
    for i, n in enumerate(a):
        positions[n].add(i)
    for (i, ai), (j, aj) in combinations(enumerate(a), 2):
        n = target - ai - aj
        if positions[n].difference((i, j)):
            return n, ai, aj
    raise NotFoundError('No triple sums to {}.'.format(target))

There are some minor optimizations you can make (for example, there is no need to record more than three positions for any given number), but these do not affect the big-O analysis, and just complicate the implementation, so I haven't made them here.

3. Comparison of algorithms

Here's a test harness for comparing the two algorithms:

from random import randint
from timeit import timeit

def test(n, m, cases=100):
    """Create cases (default: 100) random testcases and run both
    algorithms on them. Each test case consists of n numbers in the
    range [-m, m].

    """
    testcases = [[randint(-m, m) for _ in range(n)] for _ in range(cases)]
    for algorithm in triple_with_sum, triple_with_sum2:
        def run():
            for testcase in testcases:
                try:
                    assert(sum(algorithm(testcase)) == 0)
                except NotFoundError:
                    pass
        print(algorithm.__name__, timeit(run, number=1))

The triple_with_sum algorithm has the advantage that it might terminate early (after examining only a few triples), whereas triple_with_sum2 takes Ω(n) even in the best case, because it builds the dictionary of positions before starting its search. This means that triple_with_sum runs faster than triple_with_sum2 for test cases where there are many triples that sum to zero:

>>> test(1000, 100)
triple_with_sum 0.03754958091303706
triple_with_sum2 0.05412021093070507

But when triples that sum to zero are rare, the better asymptotic performance of triple_with_sum2 wins big:

>>> test(1000, 1000)
triple_with_sum 0.10221871780231595
triple_with_sum2 0.07950551761314273
>>> test(1000, 10000)
triple_with_sum 0.9003877746872604
triple_with_sum2 0.08799532195553184
>>> test(1000, 100000)
triple_with_sum 9.874092630110681
triple_with_sum2 0.14182585570961237
>>> test(1000, 1000000)
triple_with_sum 97.11457079928368
triple_with_sum2 1.0804437939077616