python - Numpy 数组多个掩码

Question

尝试根据整数掩码数组对 numpy 数组进行多次切片和平均:

即

import numpy as np

data = np.arange(11)
mask = np.array([0, 1, 1, 1, 0, 2, 2, 3, 3, 3, 3])

results = list()
for maskid in range(1,4):
    result = np.average(data[mask==maskid])
    results.append(result)
output = np.array(result)

有没有办法更快地做到这一点，也就是没有“for”循环？

Answer 1

一种方法使用 np.bincount -

np.bincount(mask, data)/np.bincount(mask)

另一个带有 np.unique 的对于一般情况，mask 中的元素不一定从 0 开始连续 -

_,ids, count = np.unique(mask, return_inverse=1, return_counts=1)
out = np.bincount(ids, data)/count