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

Question

假设我有两个数组：

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）替换为其对应的flag值也为k的所有元素的平均值？

这样转换后，x会变成：

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

（我可以使用循环来实现这个功能，但那样会非常低效。）

- Shirish Kulhari

---

那些是列表，不是数组。 - Andras Deak -- Слава Україні

@AndrasDeak：我正在使用NumPy来定义这些。 - Shirish Kulhari

下次请使用[numpy]标签。列表、stdlib数组和numpy数组都是不同的，对于你的问题来说这很重要。你的[MCVE]也应该包含数组以使其更加明显。 - Andras Deak -- Слава Україні

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 函数。它可计算分组平均值、中位数、计数、总和、最小值和最大值等统计数据。此函数还可接受用户定义的统计函数，但其性能肯定比内置统计函数慢。

- unutbu

Answer 2

一种选择是使用Pandas:

import pandas as pd
x    = [2, 4, 1, 7, 3, 9, 2, 5, 5, 1]
flag = [0, 1, 0, 2, 1, 1, 2, 0, 0, 2]
s = pd.Series(x,index=flag)
s.groupby(level=0).transform('mean').tolist()

输出：

[3.25,
 5.333333333333333,
 3.25,
 3.3333333333333335,
 5.333333333333333,
 5.333333333333333,
 3.3333333333333335,
 3.25,
 3.25,
 3.3333333333333335]

- Scott Boston

Answer 3

>>> def grouped_mean(data, flags):
...     flag_set = set(flags)
...     flags = np.asarray(flags)
...     data = np.array(data)
...     for s in flag_set:
...         m = (flags == s)
...         data[m] = np.mean(data[m])
...     return data
... 

>>> grouped_mean(x, flag)
array([ 3.25      ,  5.33333333,  3.25      ,  3.33333333,  5.33333333,
        5.33333333,  3.33333333,  3.25      ,  3.25      ,  3.33333333])

- AGN Gazer