5
\$\begingroup\$

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

\$\endgroup\$
2
  • 2
    \$\begingroup\$ Tell us about the inputs, pntl, adj_wgt and wgt_dif. Lists, arrays, dtype, shape, size, etc. \$\endgroup\$
    – hpaulj
    Jul 4 '17 at 23:20
  • \$\begingroup\$ @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. \$\endgroup\$
    – CodingNewb
    Jul 5 '17 at 17:20
2
\$\begingroup\$

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.

\$\endgroup\$
1
  • 1
    \$\begingroup\$ 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() \$\endgroup\$ Jul 5 '17 at 8:14
2
\$\begingroup\$
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.

\$\endgroup\$
4
  • \$\begingroup\$ 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? \$\endgroup\$
    – hpaulj
    Jul 4 '17 at 23:23
  • \$\begingroup\$ @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. \$\endgroup\$
    – MarianD
    Jul 5 '17 at 16:35
  • \$\begingroup\$ 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. \$\endgroup\$
    – hpaulj
    Jul 5 '17 at 16:46
  • \$\begingroup\$ @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? \$\endgroup\$
    – MarianD
    Jul 5 '17 at 17:17

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.