Element-wise addition of 2 misaligned lists?

user_10846193 0
Viewed 40

I have 2 differently sized lists (a and b) and starting index for a (a_start). I would like to add the lists together, accounting for the fact they are not aligned.

Example inputs and expected output:

a = [1,2]
b = [1,2,3,4]
a_start = 1 # 1st element of a needs to be added to 2nd element of b, etc.

expected result: [1,3,5,4]

Another example:

a = [1,2,3,4]
b = [1,2,3,4]
a_start = 2 # 1st element of a needs to be added to 3rd element of b, etc. 

expected result: [1, 2, 4, 6, 3, 4]

My current solution:

[i+j for i,j in zip([0]*a_start + a + [0]*(len(b) - len(a) - a_start), b+[0]*max(0,len(a) + a_start - len(b)))]

What is a more elegant/Pythonic way to do this for any arbitrarily sized lists?

python-3.x python
2 Answers

It would probably be best to use zip_longest:

a_padded = a_start * [0] + a
result = [x + y for x, y in itertools.zip_longest(a_padded, b, fillvalue=0)]

Also, if space is an issue, you may want to avoid creating a copy of a. In that case:

from itertools import zip_longest, chain, repeat

a_padded = chain(repeat(0, a_start), a)
result = [x + y for x, y in zip_longest(a_padded, b, fillvalue=0)]
user_3892890

Since zip_longest is already taken, here's one like yours, but I compute the length difference only once (and as a separate line, for readability) and take advantage of Python treating negative list multipliers like zero (your max(0, ...) is unnecessary):

d = a_start + len(a) - len(b)
result = [*map(sum, zip([0]*a_start + a + [0]*-d, b + [0]*d))]
user_130084390
Related