如何按mask_id快速计算二维NumPy数组的均值？

问题场景 #

我有一个二维NumPy数组，最后一列是mask_id，需要对相同mask_id对应的前三维元素计算均值：

import numpy as np

# x, y, z, mask_id
data = np.array([[ 3.28246069,  3.63800144, -1.66515279,  0],
                 [ 3.28246069,  3.63800144, -1.66515279,  0],
                 [ 0.,  0., 0.,  1]])

期望输出：

# mean value of each mask id
array([[3.28246069,  3.63800144, -1.66515279], [0, 0, 0]])

高效实现方法 #

方法1：NumPy原生unique+bincount（推荐，纯NumPy无依赖）

适合mask_id为整数的场景，无需排序，性能最优：

import numpy as np

data = np.array([[ 3.28246069,  3.63800144, -1.66515279,  0],
                 [ 3.28246069,  3.63800144, -1.66515279,  0],
                 [ 0.,  0., 0.,  1]])

# 分离数据列和mask_id列
values = data[:, :3]
mask_ids = data[:, 3].astype(int)

# 获取唯一mask_id及逆索引（用于映射原始数据到分组）
unique_ids, inverse = np.unique(mask_ids, return_inverse=True)

# 统计每个分组的元素数量
counts = np.bincount(inverse)

# 对每个维度计算加权均值（权重为1，等价于总和除以数量）
mean_values = np.array([
    np.bincount(inverse, weights=values[:, i]) / counts 
    for i in range(values.shape[1])
]).T

print(mean_values)

输出：

array([[ 3.28246069,  3.63800144, -1.66515279],
       [ 0.        ,  0.        ,  0.        ]])

方法2：NumPyreduceat（适合已排序数据）

如果数据已经按mask_id排序，此方法效率极高；未排序时需先排序：

import numpy as np

data = np.array([[ 3.28246069,  3.63800144, -1.66515279,  0],
                 [ 3.28246069,  3.63800144, -1.66515279,  0],
                 [ 0.,  0., 0.,  1]])

# 按mask_id排序
sorted_idx = np.argsort(data[:, 3])
sorted_data = data[sorted_idx]

# 获取分组分割点
unique_ids, split_points = np.unique(sorted_data[:, 3], return_index=True)

# 按分割点计算分组均值
mean_values = np.add.reduceat(sorted_data[:, :3], split_points) / np.diff(
    np.append(split_points, len(sorted_data))
)

print(mean_values)

方法3：Pandas（代码简洁，适合多操作场景）

如果允许依赖Pandas，代码可读性最强，适合同时处理其他数据操作：

import pandas as pd
import numpy as np

data = np.array([[ 3.28246069,  3.63800144, -1.66515279,  0],
                 [ 3.28246069,  3.63800144, -1.66515279,  0],
                 [ 0.,  0., 0.,  1]])

df = pd.DataFrame(data, columns=['x', 'y', 'z', 'mask_id'])
mean_values = df.groupby('mask_id')[['x', 'y', 'z']].mean().values

print(mean_values)

方法选型建议 #

• 方法1：纯NumPy实现，无额外依赖，大数据量下速度最快，优先选择。
• 方法2：适合已排序的数据，排序后计算效率接近方法1。
• 方法3：代码最简洁，适合需要同时做数据清洗、筛选等操作的场景。

内容的提问来源于stack exchange，提问作者dkanrjsk