使用NumPy的reduceat函数计算基于分组的平均值

Question

import numpy as np
import pandas as pd
dummies = np.array(pd.get_dummies(list('abdccadab'))) #categorical IV
groupIDs = np.array([10,10,10,10,20,20,30,30,30]) #groups(/strata)
_,idx,tags = np.unique(groupIDs, return_index=1, return_inverse=1)

我知道我们可以按组、列等进行求和、乘法等计算。

np.multiply.reduceat(dummies,idx)[tags]

但是有没有一种方法可以计算这些区间的平均值？

因为“np.mean.reduceat”和“np.average.reduceat”不起作用。

AttributeError: 'function' object has no attribute 'reduceat'

- Tony

---

只有 ufunc 有 reduceat。mean 和 average 不属于这种类型的函数。 - hpaulj

出于好奇，你尝试使用纯numpy的原因是什么（而不是使用pandas groupby）？ - DSM

没有特别的原因，只是我对 Pandas 比 NumPy 不够熟悉，甚至不知道在 Pandas 中也可以这样做。 - Tony

Answer 1

使用np.add.reduceat函数，基于间隔偏移量idx对数据数组dummies进行每列求和，然后使用np.bincount计算的间隔长度进行除法运算。

np.add.reduceat(dummies, idx, axis=0)/np.bincount(tags)[:,None]

计算间隔长度的另一种方法是直接使用idx-

np.diff(np.r_[idx,dummies.shape[0]])

同样地，我们可以避免使用np.unique来获得idx，如下所示 -

idx = np.r_[0,np.flatnonzero(groupIDs[1:] > groupIDs[:-1])+1]

- Divakar

---

再次感谢！我们如何将 np.add.reduceat(dummies, idx, axis=0)/np.bincount(tags)[:,None] 的输出放置到与虚拟变量相同的维度中呢？也就是说，我们将在 np.multiply.reduceat(dummies,idx) 的基础上添加 [tags] 来实现这一目标。 - Tony

@Tony，我不太清楚关于“标签”的查询。我在这里使用了tags和np.bincount一起获取每个区间的计数，从而计算平均值。不确定您打算如何与np.multiply.reduceat一起使用它。您能详细说明或重新表述一下吗？ - Divakar

抱歉表述不够清晰。当在每个函数（例如原帖中的示例）上单独调用sum，multiply等函数时，而不是通过np.bincount进行平均分割，则返回的数组与我们原始数组的形状相同。因此，我的当前问题是如何重新塑造np.add.reduceat(dummies, idx, axis=0)/np.bincount(tags)[:,None]的输出，使其具有与原始数组相同的维度。在这个例子中，第一行的前4个元素将为0.25，接下来的两个元素将为0.5等等。我正在使用所需的输出编辑我的问题。 - Tony

@Tony 如果我理解正确的话，你只需要使用 tags 进行索引：out[tags]。其中 out 是从 np.add.reduceat(dummies, idx, axis=0)/np.bincount(tags)[:,None] 得到的输出结果。这样是否有帮助？ - Divakar

绝对，先生。 - Tony

Answer 2

numpy_indexed软件包（免责声明：我是它的作者）可以提供这种类型的功能，只需一行语句即可：

import numpy_indexed as npi
unique_groups, means = npi.group_by(groupIDs).mean(dummies)

对于这种情况（已经排序的键），它提供了线性和矢量化的性能；尽管与Divakar发布的专业解决方案相比，它有更多的额外开销，因为此假设已经内置在该解决方案中。但根据您如何权衡可维护性，自我文档化和通用性，这可能是首选方案。

- Eelco Hoogendoorn