Converting 3-Sum Java to Python

I am trying to convert this Java code (I think) to Python.


    public class ThreeSum
    public static int count(int[] a)
        int N = a.length;
        int count = 0;
        for (int i = 0; i < N; i++)
            for (int j = i+1; k < N; j++)
                for (int k = j+1; k<N; k++)
                    if (a[i] + a[j] + a[k] == 0)
        return count;
    public static void main(String[] args)
        int[] a = In.readInts(args[0]);



a = [30,-40,-20,-10,40,0,10,5]
def check(list):
    N = len(list)
    count = 0
    i = 0
    while i < N:
        i += 1
        j = i+1
        while j < N:
            j += 1
            k = j+1
            while k < N:
                k += 1
                if a[i]+a[j]+a[k] == 0: #<-- this is the main part with error "list index out of range"
                    count += 1
    return count
    print count


It is supposed to check all the numbers in the list, how many sums of three pairs are equal to 0. I'm not sure what to use instead of "for (int i = 0; i <N; i ++)". I think my while loop is doing the same thing, or am I wrong?


source to share

2 answers

Python statement for

can only iterate over iteration. So, if you want to iterate over a range of numbers, you use the object range

as an iterative. So this Java:

for (int k = j+1; k<N; k++)


... becomes this Python:

for k in range(j+1, N):


It doesn't cost anything that in many cases you don't really want to do it; if the only thing you are going to use k

for is a[k]

you can just loop over more a

. However, in this case you need i

and j

to set limits j

and k

. (You can still write your loops with enumerate

, but I think it just gets in the way here.)

So, all of this can be:

N = len(a)
count = 0
for i in range(N):
    for j in range(i+1, N):
        for k in range(j+1, N):
            if a[i]+a[j]+a[k] == 0:
                count += 1
return count


However, I think it would be cleaner if you turned it into a generator transformation pipeline:

N = len(a)
sums = (a[i]+a[j]+a[k] 
        for i in range(N) for j in range(i+1, N) for k in range(j+1, N))
zeros = (n for n in sums if n == 0)
return sum(1 for _ in zeros)


Or, if you want to use the code-golf, taking advantage of the fact that True

, and False

are 1

, and 0

when used as integers:

N = len(a)
return sum(not a[i]+a[j]+a[k]
           for i in range(N) for j in range(i+1, N) for k in range(j+1, N))


But once you think in these terms, is there a way to generate all of these a[i]

, a[j]

and a[k]

in one step? Yes, these are just all possible combinations of 3 elements from a

. Thanks for @alfasin for making this clear and writing this much more succinctly than ever. :) So:

combs = itertools.combinations(a)
sums = map(sum, combs)
zerosums = (n for n in sums if n == 0)
return sum(1 for _ in zerosums)


Or, playing code golf again:

return sum(not sum(c) for c in itertools.combinations(a))




I would go with a different approach:

from itertools import combinations
a = [30,-40,-20,-10,40,0,10,5]
combs = combinations(a, 3) 
s = 0
for c in combs:
    x, y, z = c
    if x + y + z == 0:
        print x, y, z
        s += 1

print "number of combinations that sum to zero:", s



30 -40 10
30 -20 -10
-40 40 0
-10 0 10

number of combinations that sum to zero: 4




All Articles