Есть ли способ векторизации, применяя функцию среднего к замаскированным областям в ndarray?

Question

Допустим, у меня есть два ndarray, которые определены как таковые:

import numpy as np
mask = np.array([[1,1],[1,2]])
values = np.array([[1., 3.],[2., 2.]])

Моя цель - вычислить среднее значение значений на основе областей маски, указанных целым числом в mask. Естественно, я бы использовал for-l oop:

out = np.zeros(len(np.unique(mask)))
for j,i in enumerate(np.unique(mask)):
  out[j] = np.nanmean(values[mask==i])

Однако это сериализованное решение становится очень медленным для больших многомерных массивов. Есть ли способ эффективно векторизовать эту операцию? Заранее благодарим за помощь!

Answer 1

Вы можете использовать np.bincount:

unq,inv,cnt = np.unique(mask,return_inverse=1,return_counts=1)
np.bincount(inv,values.ravel())/cnt
# array([2., 2.])