# Sum array values based on conditions in other arrays

I have the following code that sums the values in wgt_dif (a numpy array) if certain conditions in two other numpy arrays are met. It is basically the equivalent of a SUMIFS function in Excel.

sum_4s = 0

for i in range(len(pntl)):
if pntl[i] == 4 and adj_wgt[i] != max_wgt:
sum_4s += wgt_dif[i]


I'm wondering if there is a more Pythonic way to write this. It works fine, but I'm new to Python and numpy and would like to expand my "vocabulary".

• Tell us about the inputs, pntl, adj_wgt and wgt_dif. Lists, arrays, dtype, shape, size, etc. Commented Jul 4, 2017 at 23:20
• @hpaulj Each of pntl, adj_wgt, and wgt_dif are of type = numpy.ndarray, shape (40,), and dtype = float64. Thank you for your thorough answer below and for introducing me to mask index arrays. Commented Jul 5, 2017 at 17:20

Let's make a test case:

In [59]: x = np.random.randint(0,10,10000)
In [60]: x.shape
Out[60]: (10000,)


(I thought test cases like this were required on Code Review. We like to have then on SO, and CR is supposed to be stricter about code completeness.)

def foo(pntl, adj_wgt, wgt_dif):
sum_4s = 0
for i in range(len(pntl)):
if pntl[i] == 4 and adj_wgt[i] != 10:
sum_4s += wgt_dif[i]
return sum_4s


Test it with lists:

In [61]: pntl = adj_wgt = wgt_dif = x.tolist() # test list versions

Out[63]: 4104
In [64]: timeit foo(pntl, adj_wgt, wgt_dif)
1000 loops, best of 3: 1.45 ms per loop


Same test with array inputs is slower (lesson - if you must loop, lists are usually better):

In [65]: timeit foo(x,x,x)
The slowest run took 5.44 times longer than the fastest. This could mean that an intermediate result is being cached.
100 loops, best of 3: 3.97 ms per loop


The suggested list comprehension is modestly faster

In [66]: sum([w for w, p, a in zip(wgt_dif, pntl, adj_wgt) if p == 4 and a != 10])
Out[66]: 4104
In [67]: timeit sum([w for w, p, a in zip(wgt_dif, pntl, adj_wgt) if p == 4 and a != 10])
1000 loops, best of 3: 1.14 ms per loop


foo could have been written with zip instead of the indexed iteration. (todo - time that).

But since you say these are arrays, let's try a numpy version:

def foon(pntl, adj_wgt, wgt_dif):
# array version

In [69]: foon(x,x,x)
Out[69]: 4104
In [70]: timeit foon(x,x,x)
10000 loops, best of 3: 105 µs per loop


This is an order of magnitude faster. So if you already have arrays, try to work with them directly, without iteration.

def foo2(pntl, adj_wgt, wgt_dif):
sum_4s = 0
for w, p, a in zip(wgt_dif, pntl, adj_wgt):
if p == 4 and a != 10:
sum_4s += w
return sum_4s
Out[77]: 4104
In [78]: timeit foo2(pntl, adj_wgt, wgt_dif)
1000 loops, best of 3: 1.17 ms per loop


So it's the zip that speeds up your original code, not the list comprehension.

• For the OP: the important bit in this answer is the use of a mask index array: wgt_dif[(pntl == 4) & (adj_wgt != 10)].sum() Commented Jul 5, 2017 at 8:14
sum([w for w, p, a in zip(wgt_dif, pntl, adj_wgt) if p == 4 and a != max_wgt])


Explanation:

zip(a, b, c)


creates the list of triplets of corresponding values from the lists a, b, c - something as

[(a[0], b[0], c[0]), (a[1], b[1], c[1]), (a[2], b[2], c[2]), ...]


so the part

for w, p, a in zip(wgt_dif, pntl, adj_wgt)


loops over this triples, associating th 1st item to w, 2nd to p, and 3rd to a.

• And how does this compare to the OP's for loop version? Other than being a one liner, it doesn't look any simpler or clearer. Is it faster? Commented Jul 4, 2017 at 23:23
• @hpaulj - "I'm wondering if there is a more Pythonic way to write this" was the only OP question. Is something bad with the fact that I answered it? BTW, in my answer IS something new (sum()) and it is both simpler and cleaner - and I'm sure it will be for you, too sometime in future. Commented Jul 5, 2017 at 16:35
• I showed in my answer that your approach provides a modest speed increase, mainly from the use of zip. In simple cases list comprehensions also improve speed. But as they get more complex they lose both the speed and clarity advantage. Commented Jul 5, 2017 at 16:46
• @hpaulj - Your answer is really very nice one - in spite of you didn't answer the OP question, I'm sorry. It was not about the speed. My list comprehension is not a complex one, is it? Commented Jul 5, 2017 at 17:17