addition of two list elementwise with abortion if condition satisfied python

Question

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:
            return
        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 ?

Answer 1

You're operating on individual "bits" with Python code. Can be faster to use data types that do the operations in C. One way is to go array→bytes→int→bytes→array, doing the operations at the int level:

def dont_talk_just_code(lis1, lis2):
    b1 = lis1.tobytes()
    b2 = lis2.tobytes()
    i1 = int.from_bytes(b1, 'big')
    i2 = int.from_bytes(b2, 'big')
    if i1 & i2:
        return
    b = (i1 | i2).to_bytes(len(b1), 'big')
    return array('H', b)

Benchmark results when there is no sum larger than 1:

18.77 μs  18.97 μs  19.23 μs  original
19.24 μs  19.39 μs  19.57 μs  _8349697
 1.74 μs   1.74 μs   1.78 μs  dont_talk_just_code

Benchmark results where there's a sum larger than 1 at index 50:

 7.94 μs   7.98 μs   8.02 μs  original
 2.48 μs   2.48 μs   2.52 μs  _8349697
 1.10 μs   1.11 μs   1.11 μs  dont_talk_just_code

Maybe array→bytes→numpy→bytes→array would be even faster, I didn't try.

Benchmark code (Try it online!):

from timeit import repeat
import random
from array import array
from functools import partial

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

def _8349697(lis1, lis2):
    if (1, 1) in zip(lis1, lis2):
        return
    return array('H', (i + j for i, j in zip(lis1, lis2)))

def dont_talk_just_code(lis1, lis2):
    b1 = lis1.tobytes()
    b2 = lis2.tobytes()
    i1 = int.from_bytes(b1, 'big')
    i2 = int.from_bytes(b2, 'big')
    if i1 & i2:
        return
    b = (i1 | i2).to_bytes(len(b1), 'big')
    return array('H', b)

funcs = original, _8349697, dont_talk_just_code

def create_valid():
    added = random.choices([(0, 0), (0, 1), (1, 0)], k=100)
    lis1, lis2 = zip(*added)
    lis1 = array('H', lis1)
    lis2 = array('H', lis2)
    return lis1, lis2

for _ in range(3):
    lis1, lis2 = create_valid()
    # lis1[50] = lis2[50] = 1
    for func in funcs:
        # print(func(lis1, lis2))
        number = 10000
        times = sorted(repeat(partial(func, lis1, lis2), number=number))[:3]
        print(*('%5.2f μs ' % (t / number * 1e6) for t in times), func.__name__)
    print()

Answer 2

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