如何用Numpy高效批量将一维数组转为含字典的二维数组？

问题描述 #

现有两个形状为(N)的Numpy数组，需要生成一个形状为(N, 1)的新数组，其中每个元素是一个字典，键对应数组名称，值为对应索引的数组元素。

原实现方式效率极低（尤其当N很大时）：

import numpy as np

a = np.array([1, 2, 3, 4])
b = np.array([-1, -2, -3, -4])

result = []
for i in range(a.shape[0]):
    result.append(np.array([{"val_a": a[i], "val_b": b[i]}]))
result = np.array(result)

print(result.shape)
print(result)

输出结果：

(4, 1)
[[{'val_a': 1, 'val_b': -1}]
 [{'val_a': 2, 'val_b': -2}]
 [{'val_a': 3, 'val_b': -3}]
 [{'val_a': 4, 'val_b': -4}]]

高效实现方案 #

由于Numpy对Python对象（比如字典）的存储本质是object类型数组，无法做到完全的向量化运算，但可以通过批量生成字典列表再一次性转数组的方式大幅提升效率：

import numpy as np

a = np.array([1, 2, 3, 4])
b = np.array([-1, -2, -3, -4])

# 用列表推导批量生成字典，再转为(N,1)形状的数组
dict_list = [{"val_a": x, "val_b": y} for x, y in zip(a, b)]
result = np.array(dict_list).reshape(-1, 1)

print(result.shape)
print(result)

效率对比说明 #

• 原代码的问题在于：循环中每次都创建一个单元素Numpy数组并追加，频繁的数组对象初始化和内存分配会带来巨大的性能开销，N越大越明显。
• 改进后的方式：利用Python内置的列表推导（高度优化的循环实现）批量生成所有字典，再一次性转换为Numpy数组并调整形状，避免了重复的数组创建操作，在N很大时性能提升非常显著。

如果你的场景允许使用结构化数组而非字典元素，也可以考虑Numpy的结构化数组方案（更贴合Numpy的设计，性能更优），示例如下：

# 结构化数组方案（元素为结构化对象，而非字典）
structured_arr = np.array(list(zip(a, b)), dtype=[('val_a', int), ('val_b', int)])
# 转为(N,1)形状
structured_arr = structured_arr.reshape(-1, 1)
print(structured_arr.shape)
print(structured_arr)

输出：

(4, 1)
[[( 1, -1)]
 [( 2, -2)]
 [( 3, -3)]
 [( 4, -4)]]

这种方式的元素可以通过structured_arr[i,0]['val_a']访问，性能比存储字典的object数组更高，适合对性能要求极高的场景。

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