Apc
Apc

Reputation: 25

indexing one array and doing operation on another

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

Answers (1)

Kasravnd
Kasravnd

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

Related Questions