I have a function that sums two arrays element-wise and stops if the sum at a certain index is greater strictly than 1. This function is called lot of times (15817145 times for a small simulation and expects to grow a lot more as the simulation size grows). Profiling my code shows that it is among the five most consuming functions of my code and I believe can be optimized further.

def add_lis(lis1, lis2):
    res = []
    for i, j in zip(lis1, lis2):
        if i + j > 1:
        res.append(i + j)
    return array('H', res)

lis1 and lis2 are both of type array.array('H',..). The two lists have the same length which can be up to 100 and they are both filled with zeroes and ones only. I tried with numpy arrays but there is too much overhead that it is not beneficial to use it.

My question is can this function be optimized? Maybe in a sort of list comprehension ?

  • 2
    \$\begingroup\$ You have framed this as a question of speed. But you also say the lists have 100 or fewer elements (ie, tiny) and the function is called several times (meaning unclear, but it also sounds small-scale). Have you benchmarked the code? How much time is this function costing you and is it worth the effort to optimize? Alternatively, perhaps this isn't really a speed issue and you are curious about different ways someone might write this code more clearly or compactly. Perhaps you can edit your question to clarify your goal or at least explain how a substantive speed issue is at play. \$\endgroup\$ – FMc May 12 at 18:38
  • \$\begingroup\$ If lis1 and lis2 are long, finding the index of the first above-one sum will be fast with numpy, and this function can be vectorized. But you need to show typical calls to this function, and what lis1 and lis2 really contain, along with expected lengths. \$\endgroup\$ – Reinderien May 12 at 19:03
  • 1
    \$\begingroup\$ hope my edits clear up your questions @FMc \$\endgroup\$ – sos May 14 at 6:10
  • \$\begingroup\$ @sos Yes, indeed, that clears things up. I think the suggestion to try numpy is a good one. It's a library designed to do such things more quickly at a large scale like that -- at least that's my rough understanding of it. \$\endgroup\$ – FMc May 14 at 14:59

Does it make sense (in the context of your program execution) to split your task into two stages?

Check for the presence of (1, 1) in zip(lis1, lis2), and if False, then pass the generator to the array: array('H', (i + j for i, j in zip(lis1, lis2))).


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.