Converting 3-Sum Java to Python
I am trying to convert this Java code (I think) to Python.
From:
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)
count++;
return count;
}
public static void main(String[] args)
{
int[] a = In.readInts(args[0]);
StdOut.printIn(count(a));
}
}
To:
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
check(a)
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
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))
source to share
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
OUTPUT
30 -40 10
30 -20 -10
-40 40 0
-10 0 10
number of combinations that sum to zero: 4
source to share