Python:将分组均值分配给1-D数组

Question

假设我有2个数组:

x    = [2, 4, 1, 7, 3, 9, 2, 5, 5, 1]
flag = [0, 1, 0, 2, 1, 1, 2, 0, 0, 2]

该flag数组表示每个元素x属于哪个"组" .如何x用相应值k的所有元素的平均值替换(例如,标记值)的每个元素？xflagk

经过这样的改造,x看起来像:

x    = [3.25, 5.33, 3.25, 3.33, 5.33, 5.33, 3.33, 3.25, 3.25, 3.33]

(我可以使用循环来实现这一点,但效率很低.)

Answer 1

您可以np.bincount用来计算分组方法:

import numpy as np
x    = np.array([2, 4, 1, 7, 3, 9, 2, 5, 5, 1])
flag = np.array([0, 1, 0, 2, 1, 1, 2, 0, 0, 2])
total = np.bincount(flag, weights=x)
count = np.bincount(flag)
means = (total/count)[flag]

产量

array([ 3.25      ,  5.33333333,  3.25      ,  3.33333333,  5.33333333,
        5.33333333,  3.33333333,  3.25      ,  3.25      ,  3.33333333])

对于更通用的分组统计,还有scipy.stats.binned_statistic 功能.它可以计算分组平均值,中位数,计数,总和,最小值,最大值统计量.它还可以接受统计信息的用户定义函数,但性能(当然)比内置统计信息慢.