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

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?

CodePudding user response：

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