使用Numpy的unique函数后将关联id按去重行分组的实现方法

核心是利用np.unique的return_inverse参数，该参数会返回原始数组中每一行对应的唯一结果行的下标，通过这个下标即可快速完成ID分组。

完整实现代码如下：

import numpy as np

# 你的输入数据
t = np.array([[0, 3, 4], [1, 2, 8], [1, 2, 8]])
ids = np.array(['A', 'B', 'C'])

# 调用unique时开启return_inverse返回逆索引
unique_rows, inverse_indices = np.unique(t, axis=0, return_inverse=True)

# 按逆索引分组ID，生成长度不一致的object类型numpy数组
grouped_ids = np.array([ids[inverse_indices == i] for i in range(len(unique_rows))], dtype=object)

输出验证：

print(unique_rows)
# 输出：
# [[0 3 4]
#  [1 2 8]]

print(grouped_ids)
# 输出：
# [array(['A'], dtype='<U1') array(['B', 'C'], dtype='<U1')]

如果你不需要严格返回numpy数组，只需要对应结构的列表，可调整最后一步的写法：

grouped_ids = [ids[inverse_indices == i].tolist() for i in range(len(unique_rows))]
# 输出结果为 [['A'], ['B', 'C']]，和你预期的结构完全一致

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