Sum array values based on conditions in other arrays

Question

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".

Answer 1

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.)

Your code as a function:

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

In [63]: foo(pntl, adj_wgt, wgt_dif)
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
    mask = (pntl==4) & (adj_wgt != 10)
    return wgt_dif[mask].sum()

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
In [77]: foo2(pntl, adj_wgt, wgt_dif)
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.

Answer 2

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 1^st item to w, 2^nd to p, and 3^rd to a.