CODEWITHSUNDEEP
pythonnumpy
CR7
CR7

Reputation: 135

numpy - one array calculates the average based on the position of another array

I have 2 arrays with the same shape. If the value of the element of the bList array corresponding to the aList array is 255, then find the corresponding position in the aList array, and add the eligible elements of the a array to calculate the average.

I think I can do it with loop but I think it's stupid.

import numpy as np

aList = np.array([[1,2,3,4], [5,6,7,8], [9,10,11,12]])
bList = np.array([[255,255,0,255], [0,0,255,255], [255,0,0,0]])
sum_list = []
for id,row in enumerate(aList):
    for index,ele in enumerate(row):
        if bList[id][index]==255:
            tmp = aList[id][index]
            sum_list.append(tmp)
average = np.mean(sum_list) #(1+2+4+7+8+9)/6 #5.166666666666667

Is there a simpler way?

Upvotes: 1

Views: 650

Answers (1)

mozway
mozway

Reputation: 260640

Use numpy.where

np.mean(aList[np.where(bList==255)])

Or with a boolean mask:

mask = (bList==255)

(aList*mask).sum()/mask.sum()

Output: 5.166666666666667

Upvotes: 1

Related Questions