基于索引转换字典的Numpy广播优化：寻求替代循环的向量化实现方案

当然可以用向量化操作来替代for循环，这在处理大规模稀疏数组时能大幅提升效率！下面我会给出两种高效的实现方案，完全贴合你的需求：

方案一：基于有效行的映射拼接（适合稀疏场景） #

这个方案针对arr1稀疏的特点，只处理非NaN的行，避免不必要的计算：

import numpy as np

# 你的原始数据
d = {0:[0,1,3], 1:[4,5,6], 2:[2,7,8], 3:[9]}
arr1 = np.array([[0,1], [np.NaN,np.NaN], [np.NaN,6.5], [16,0.2]])

# 1. 筛选arr1中不全为NaN的有效行索引
valid_indices = np.where(~np.isnan(arr1).all(axis=1))[0]

# 2. 构建arr2的目标索引数组：把每个有效索引对应的d中的列表拼接起来
target_indices = np.concatenate([d[idx] for idx in valid_indices])

# 3. 把arr1的有效行重复对应次数（和d中列表长度匹配）
repeat_counts = [len(d[idx]) for idx in valid_indices]
arr1_valid_repeated = np.repeat(arr1[valid_indices], repeat_counts, axis=0)

# 4. 初始化arr2并赋值
arr2_shape = (sum(len(v) for v in d.values()), arr1.shape[1])
arr2 = np.full(arr2_shape, np.nan)
arr2[target_indices] = arr1_valid_repeated

逻辑验证

运行这段代码后，arr2的结果和你原循环的输出完全一致：

• 索引0、1、3的行是[0,1]
• 索引2、7、8的行是[np.NaN,6.5]
• 索引9的行是[16,0.2]
• 其余行保持默认的np.NaN

方案二：全局映射数组赋值（适合字典键数量少的场景） #

如果字典d的键数量不多（比如你的场景是600个），可以先构建一个从arr2索引到arr1索引的映射数组，再批量赋值：

import numpy as np

# 你的原始数据
d = {0:[0,1,3], 1:[4,5,6], 2:[2,7,8], 3:[9]}
arr1 = np.array([[0,1], [np.NaN,np.NaN], [np.NaN,6.5], [16,0.2]])

# 1. 构建映射数组：map_arr[i]表示arr2的第i行对应arr1的第map_arr[i]行
arr2_len = sum(len(v) for v in d.values())
map_arr = np.full(arr2_len, -1)  # 用-1标记无对应关系的位置
for idx, target_list in d.items():
    map_arr[target_list] = idx

# 2. 批量赋值：直接用映射数组索引arr1
arr2 = arr1[map_arr]

# 3. 处理arr1中原本为NaN的行对应的位置
nan_row_mask = np.isnan(arr1).all(axis=1)[map_arr]
arr2[nan_row_mask] = np.nan

为什么这两种方案更快？ #

你的原for循环是Python级别的迭代，每次循环都要处理一行数据，在arr2规模达到19万行时，效率会很低。而上面的方案用的是numpy的向量化操作，这些操作是基于C语言底层实现的，能充分利用CPU的并行计算能力，处理大规模数据时速度会提升数倍甚至数十倍。

对于你的场景（arr1是600行稀疏数组，arr2是19万行），方案一的效率会更高，因为它只处理arr1中的有效行，避免了对全NaN行的无用处理。

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