Reputation: 25
I got two arrays as
a = [0, 1, 2, 0, 0, 0, 1, 2, 3, 0, 0, 1, 2]
b = [1, 2, 1, 2, 3, 1, 4, 5, 1, 5, 6, 7, 8]
Indexing on first array a
, I want to do some calculation on b
as:
sum the elements of b
from one zero on a
upto the arrival of next zero, but excluding the elements of b
at next zero. And, if there are two consecutive zeros on a
same elements should be stored in new vector containing the result.
Final vector should be:
result=[4,2,3,11,5,21].
I tried like this:
vec1 = []
for i in range(len(a)-1):
if a[i] == a[i+1] == 0:
vec1.append(b[i])
print(vec1)
But I got the result only for consecutive zeros. I couldn't figure out for the case when there is no consecutive zeros.
Upvotes: 1
Views: 37
Reputation: 107287
Firs of all you need to find the index of those element in a that are 0,then zip a
and b
to create a pair of those elements, and then with a list comprehension find grub the relevant sum but pay attention to demo for better understanding the code :
>>> a=[0,1,2,0,0,0,1,2,3,0,0,1,2]
>>> b=[1,2,1,2,3,1,4,5,1,5,6,7,8]
>>> indx=[i for i,j in enumerate(a) if j==0]+[len(a)]
>>> l= list(zip(a,b))
>>> [sum(v[1] for v in l[i:j]) for i,j in zip(indx,indx[1:])]
[4, 2, 3, 11, 5, 21]
Demo:
>>> indx=[i for i,j in enumerate(a) if j==0]+[len(a)]
>>> indx
[0, 3, 4, 5, 9, 10, 13]
>>> zip(indx,indx[1:])
[(0, 3), (3, 4), (4, 5), (5, 9), (9, 10), (10, 13)]
Upvotes: 1