Numpy按前两列分组并计算组内第三列均值的高效实现咨询

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

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

我的需求是：

1. 按前两列对数组分组，得到类似grouped_arr = np.array([[[0,1,1], [0,1,2]], [[2,2,2]]])的分组结果
2. 每组仅保留一条记录，第三列为该组第三列的均值，最终得到：

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

我尝试了以下代码，但不确定其正确性与执行效率（实际数据类型为float）：

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.])]

你的代码逻辑是正确的，但嵌套循环会在数据量较大时拖慢效率——Numpy的核心优势是矢量化操作，而非逐元素循环。下面提供两种更可靠高效的实现方式：

方法一：纯Numpy矢量化实现 #

利用np.unique的return_inverse参数获取分组索引，再结合np.bincount批量计算均值：

import numpy as np

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

# 获取前两列的唯一值及对应分组索引
unique_keys, idx = np.unique(arr[:, :2], axis=0, return_inverse=True)

# 计算每个分组的第三列均值：总和/数量
counts = np.bincount(idx)
sums = np.bincount(idx, weights=arr[:, 2])
means = sums / counts

# 拼接唯一键和均值得到最终结果
final_array = np.column_stack((unique_keys, means))
print(final_array)
# 输出：
# [[0.  1.  1.5]
#  [2.  2.  2. ]]

这种方法完全基于矢量化操作，避免了循环，在大数据量下效率远高于原始代码。

方法二：使用Pandas（适合复杂分组场景） #

如果数据规模很大，或者需要更灵活的分组逻辑，Pandas的分组功能经过优化，效率和可读性都更出色：

import numpy as np
import pandas as pd

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

# 转换为DataFrame
df = pd.DataFrame(arr, columns=['col1', 'col2', 'col3'])
# 按前两列分组，计算第三列均值并重置索引
result_df = df.groupby(['col1', 'col2'])['col3'].mean().reset_index()
# 转换回Numpy数组
final_array = result_df.to_numpy()

print(final_array)
# 输出：
# [[0.  1.  1.5]
#  [2.  2.  2. ]]

原始代码的效率瓶颈 #

你的代码存在两个主要问题：

• 循环遍历唯一值时，每次arr[v]都会执行一次布尔索引，数据量大时会重复创建数组，开销很高
• 嵌套循环没有利用Numpy的矢量化特性，逐组计算均值的速度远慢于批量计算

上述两种方法都通过批量处理避免了这些问题，同时保证结果正确。

内容的提问来源于stack exchange，提问作者Below the Radar