遍历行长度不一的NumPy数组 生成目标数组c的技术需求

问题描述 #

给定10行5列的数组a和20个元素的数组b（原始数据如下），需要将a中的每个非零元素n替换为b中第n个元素，a中的0元素保持为0，最终得到10x5的数组c。

原始数组：

# 注：原数组a第8行存在格式错误，已修正为5列
a = [[15,2,1,0,0],
     [8,12,0,0,0],
     [4,12,10,9,0],
     [3,0,0,0,0],
     [19,7,0,0,0],
     [13,15,4,0,0],
     [0,0,0,0,0],
     [11,4,0,0,0],
     [7,0,0,0,0],
     [10,6,8,4,0]]

b = [8,1,6,4,9,3,5,6,11,14,5,4,33,7,9,15,7,3,19,3]

需求示例：数组c的前两行应为：

[[9,1,8,0,0],
 [6,4,0,0,0]]

实现代码 #

import numpy as np

# 修正并转换为NumPy数组
a = np.array([
    [15,2,1,0,0],
    [8,12,0,0,0],
    [4,12,10,9,0],
    [3,0,0,0,0],
    [19,7,0,0,0],
    [13,15,4,0,0],
    [0,0,0,0,0],
    [11,4,0,0,0],
    [7,0,0,0,0],
    [10,6,8,4,0]
])

b = np.array([8,1,6,4,9,3,5,6,11,14,5,4,33,7,9,15,7,3,19,3])

# 初始化全0结果数组
c = np.zeros_like(a, dtype=int)

# 定位非零元素位置并完成映射
non_zero_mask = a != 0
c[non_zero_mask] = b[a[non_zero_mask] - 1]

# 输出结果
print(c)

代码说明 #

1. 数组修正与转换：将原始列表转为NumPy数组，同时修正a中第8行的格式错误（原数据多一个0，调整为5列以匹配10x5的形状要求）。
2. 结果数组初始化：用np.zeros_like创建与a形状、数据类型完全一致的全0数组c，确保后续赋值格式匹配。
3. 矢量化映射：通过布尔掩码non_zero_mask快速定位a中所有非零元素的位置，利用NumPy矢量化索引特性，将a的非零元素值减1（适配Python/NumPy的0索引规则）后作为b的索引，取出对应值直接赋值给c的对应位置。这种方式比逐行循环效率更高，适合大规模数组处理。

验证结果 #

运行代码后，输出的前两行与需求示例完全一致，整个数组c的结果符合要求。

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