使用不规则的索引列表对数组进行索引，并执行求和/平均值缩减

Question

我有一些2D数据，其中第一个轴是时间，第二个轴是人的ID。因此，数据条目是人随时间的属性值。

我想做的是对人员进行分组，并在所有时间段内对每组中的属性进行平均。以下是6个时间点和5个人的样本，其中有2组

import numpy as np
data = np.arange(30)
data.shape = 6, 5
groups = [[0, 1, 4], [2, 3]]
result = np.empty((6, 2))
for i, indices in enumerate(groups):
result[:, i] = data[:, indices].mean(axis=1)

result是

array([[ 1.66666667,  2.5       ],
[ 6.66666667,  7.5       ],
[11.66666667, 12.5       ],
[16.66666667, 17.5       ],
[21.66666667, 22.5       ],
[26.66666667, 27.5       ]])

就效率而言，这是我们能做的最好的事情吗？我想知道是否也可以消除groups上的循环。

Answer 1

方法#1：一般情况

这是一种利用np.add.reduceat-的几乎矢量化的方法

g = np.concatenate(groups)
lens = list(map(len, groups))
cuts = np.r_[0,np.cumsum(lens)[:-1]]
out = np.add.reduceat(data[:, g], cuts, axis=1)/lens

方法#2：具体情况

如果groups是一个常规的2D数组/列表，我们可以简单地做-

data[:, groups].mean(2)

方法#3：方法1+2

混合方法1和2，我们可以提出另一种通用的案例方法-

from itertools import zip_longest
idx = np.vstack(list(zip_longest(*groups, fillvalue=-1)))
c = (idx == -1).sum(0)
sums = data[:, idx].sum(1) - c*(data[:,-1,None])
lens = list(map(len, groups))
out = sums/lens

方法#4：使用矩阵乘法

我们将创建一个掩码，当矩阵与data相乘时，其求和约简将给我们求和约简的版本，然后除以长度，得到我们想要的平均值-

mask = np.zeros((data.shape[1], len(groups)), dtype=bool)
for i, indices in enumerate(groups):
mask[indices,i] = 1
out = data.dot(mask)/list(map(len, groups))

此外，我们可能希望使用float32来实现更快的矩阵乘法-

data.dot(mask.astype(np.float32))

方法#5：方法2+3

我们将使用零作为附加列，并使用由zip_longest创建的常规索引数组，对其进行索引和求和，从而获得平均值-

data0 = np.pad(data,((0,0),(0,1)))
idx = np.vstack(list(zip_longest(*groups, fillvalue=-1)))
out = data0[:, idx].sum(1)/list(map(len, groups))