python - Elegant way to complete a parallel operation on two arrays of unequal lengths -

i want write function in numba runs math operation on 2 arrays, , accommodate when both arrays don't have same element count.

so example: lets want function adds each element of array a elements of array b these 3 possible scenarios:

1) both a , b have same number of items, c[ii]=a[ii]+b[ii]

2) a has more items b: c[ii]=a[ii]+b[ii] until b's upper limit, , complete c[ii]=a[ii]+b[-1]

3) a has fewer items b: c[ii]=a[ii]+b[ii] until a's upper limit, , complete c[ii]=a[-1]+b[ii]

for wrote code below, works fine , fast when dealing millions of values, can see 3 identical code blocks feel terribly wasteful. plus if/else running in loop feels terrible.

from numba import jit, float64, int32  @jit(float64[:](float64[:], float64[:]), nopython=true) def add(a, b):      # both shapes equal: add between a[i] , b[i]     if a.shape[0] == b.shape[0]:         c = np.empty(a.shape)          in range(a.shape[0]):             c[i] = a[i] + b[i]          return c      # has more entries b: add between a[i] , b[i] until b.shape[0]-1 reached.     # finish loop add between a[i] , b[-1]     elif a.shape[0] > b.shape[0]:         c = np.empty(a.shape)         i_ = b.shape[0]-1 # upper limit of b's shape          in range(a.shape[0]):             if < b.shape[0]:                 c[i] = a[i] + b[i]             else:                 c[i] = a[i] + b[i_]          return c      # b has more entries a: add between a[i] , b[i] until a.shape[0]-1 reached.     # finish loop add between a[-1] , b[i]         else:         c = np.empty(b.shape)         i_ = a.shape[0]-1 # upper limit of a's shape          in range(b.shape[0]):             if < a.shape[0]:                 c[i] = a[i] + b[i]             else:                 c[i] = a[i_] + b[i]          return c 

i'm new numba , jit python code compilation may "the efficient way" want.

but if there more elegant way without sacrificing speed love know how.

but can see 3 identical code blocks feel terribly wasteful.

yes, you're repeating lot in code. on other hand it's easy see each case does.

you use 2 loops instead:

import numba nb  @nb.njit(nb.float64[:](nb.float64[:], nb.float64[:])) def add2(a, b):     size1, size2 = a.shape[0], b.shape[0]     maxsize, minsize = max(size1, size2), min(size1, size2)     c = np.empty(maxsize)      # calculate elements present in , b     idx in range(minsize):         c[idx] = a[idx] + b[idx]      # check array longer , fillvalue should applied     if size1 > size2:         missing =         filler = b[-1]     else:         missing = b         filler = a[-1]      # calculate elements after or b ended. in case have equal lengths     # range of length 0 won't enter.     idx in range(minsize, maxsize):         c[idx] = missing[idx] + filler      return c 

a lot less repetition, maybe not clear.

plus if/else running in loop feels terrible.

actually it's not bad seems because branch prediction makes if cheap. true long both arrays still have elements , switches false (and stays false thereafter) when 1 array exhausted. that's easy computer predict check cheap (almost free).