如何将三列Numpy数组转换为元组列表构成的字典？

给定如下三列Numpy数组：

import numpy as np
arr = np.array([[1,2,3,3],[5,6,7,8],[8,7,6,5]]).T

# 数组内容：
# array([[1, 5, 8],
#        [2, 6, 7],
#        [3, 7, 6],
#        [3, 8, 5]])

需要将其转换为以第一列为键、后两列元组列表为值的字典，目标格式：

{1:[(5,8)], 2:[(6,7)],3: [(7,6),(8,5)]}

方法一：Numpy向量化分组（大数据量首选） #

利用Numpy的unique方法获取分组索引，批量提取对应元素，避免Python循环的性能损耗，适合处理大规模数组：

import numpy as np

arr = np.array([[1,2,3,3],[5,6,7,8],[8,7,6,5]]).T

# 获取第一列的唯一键值
keys = np.unique(arr[:, 0])

# 遍历唯一键，提取对应行的后两列并转为元组列表
result = {}
for key in keys:
    target_rows = arr[arr[:, 0] == key, 1:]
    result[key] = [tuple(row) for row in target_rows]

print(result)

方法二：defaultdict循环（小数据量简洁方案） #

如果数组规模不大，用collections.defaultdict配合循环实现，代码更直观简洁：

import numpy as np
from collections import defaultdict

arr = np.array([[1,2,3,3],[5,6,7,8],[8,7,6,5]]).T

result = defaultdict(list)
for row in arr:
    result[row[0]].append(tuple(row[1:]))

# 可选：转换为普通字典
result = dict(result)
print(result)

性能说明 #

• 当数组行数超过10万级时，方法一的Numpy向量化操作速度远快于纯Python循环；
• 小数据量场景下，两种方法性能差异可忽略，方法二的代码可读性更强。

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