如何快速并行处理2×N列表/NumPy数组，按首行同值分组计算第二行数据

以下两种方案都避免了Python层面的显式遍历，底层依赖优化过的C语言实现运算，运行效率远高于手动循环实现，适合大规模数据处理场景：

纯NumPy向量化方案 #

全程使用NumPy原生向量化操作，没有额外依赖，核心利用排序+分组边界识别+reduceat批量聚合实现：

import numpy as np

# 输入示例数组
a = np.array([[3,   2,   1,   2,   3,   3  ],
              [3.1, 2.2, 1.1, 2.1, 3.3, 3.2]])

# 按第一行值排序，将同组元素集中
sort_idx = np.argsort(a[0])
sorted_key = a[0][sort_idx]
sorted_val = a[1][sort_idx]

# 识别分组边界
unique_keys, split_pos = np.unique(sorted_key, return_index=True)

# 批量计算分组求和、分组元素数量、分组均值
group_sum = np.add.reduceat(sorted_val, split_pos)
group_count = np.diff(np.append(split_pos, len(sorted_val)))
group_mean = group_sum / group_count

# 批量计算每个元素与所属组均值的绝对差，再按分组求和
mean_expanded = np.repeat(group_mean, group_count)
abs_deviation = np.abs(sorted_val - mean_expanded)
result = np.add.reduceat(abs_deviation, split_pos)

运行后得到的result和你原有实现输出完全一致。

Pandas分组聚合方案 #

如果允许引入Pandas依赖，代码会更简洁，Pandas的groupby底层做了大量性能优化，处理多分组、大数据量场景时表现优异：

import pandas as pd
import numpy as np

a = np.array([[3,   2,   1,   2,   3,   3  ],
              [3.1, 2.2, 1.1, 2.1, 3.3, 3.2]])

df = pd.DataFrame(a.T, columns=["group_key", "value"])
result = df.groupby("group_key")["value"].apply(
    lambda x: np.sum(np.abs(x - x.mean()))
).values

性能说明 #

你原有实现的时间复杂度为O(n*k)，其中n是数组长度、k是分组数量，每次循环都要做全量布尔匹配，数据量大时开销极高。上述两种方案的主要开销来自排序的O(n log n)，后续聚合操作都是O(n)线性复杂度，数据量越大，性能优势越明显。

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