Numpy - 计算组的平均值

Question

例如，我有一个形状为 (N, 3) 的数组：

arr = np.array(
  [
    [0,1,1], 
    [0,1,2], 
    [2,2,2]
  ]
)

我想按前两列对这个数组进行分组以获得这样的数组

grouped_arr = np.array([[[0,1,1], [0,1,2]], [[2,2,2]]])

最后我想按组只得到一个元素，而第三列将是组第三列的平均值

final_array = np.array([[0,1,1.5], [2,2,2]])

我正在尝试一些东西但不确定它是否正确以及它是否是实现它的有效方法：

import numpy as np

arr = np.array([[0,1,1], [0,1,2], [2,2,2]])

stacked = np.vstack((arr[:,0], arr[:,1])).transpose()
uniques_values = np.unique(stacked, axis=0)

groups = []
for v in uniques_values:
    groups.append(arr[v])

final_arr = []
for group in groups:
    mean = np.mean(group[:,2], axis=0)
    final_arr.append(np.array([group[0][0], group[0][1], mean]))

print(final_arr)

>>>[array([0. , 1. , 1.5]), array([2., 2., 2.])]

我正在寻找可靠且有效的建议。在我的真实数据中，dtype 是 float

Answer 1

您的“分组”实际上只是普通切片：

import numpy as np

arr = np.array((
    (0,1,1),
    (0,1,2),
    (2,2,2),
))

final_array = np.vstack((
    arr[:-1].mean(axis=0),
    arr[-1],
))

print(final_array)

[[0.  1.  1.5]
 [2.  2.  2. ]]

这将针对您规定的形状 3,3 的输入数据运行并产生您规定的输出。

Answer 2

在 numpy 中不能有参差不齐的数组。这就是列表的用途。您可以识别组的位置，并使用 ufuncs 提供的聚合函数来操作它们，而不是将数据分成组。

具体来说，一旦确定了组开始的索引，就可以使用

np.add.reduceat

来构建均值。当前两列中的任一数字与上一行中的数字不匹配时，组开始：

mask = (arr[:-1, :2] != arr[1:, :2]).any(1)
group_indices = np.flatnonzero(np.r_[True, mask])

现在取平均值很简单：

sums = np.add.reduceat(arr[:, -1], group_indices, axis=0)
lengths = np.diff(np.r_[group_indices, len(arr)])
means = sums / lengths

要构造输出，您只需要数组的前两列在每个组的起始索引处与均值连接（转置为一列）：

result = np.concatenate((arr[group_indices, :2], means[:, None]), axis=1)