HiveBrain v1.2.0
Get Started
← Back to all entries
patternpythonModerate

Given an array find any three numbers which sum to zero

Submitted by: @import:stackexchange-codereview··
0
Viewed 0 times
threesumarrayanynumbersfindwhichzerogiven

Problem

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

Solution

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 []

Code Snippets

def test_winston1(self):
    res = sum_to_zero([125, 124, -100, -25])
    self.assertEqual(res, [125,-100,-25])
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):
for k in xrange(j + 1, len_a):
            if a[j] + a[k] == -1 * a[i]:
return [a[i], a[j], a[k]]
    return []

Context

StackExchange Code Review Q#23996, answer score: 15

Revisions (0)

No revisions yet.