如何将ndarray大部分元素置零，仅保留指定索引位置的数值

针对你的需求，不需要逐行循环，使用Numpy原生高级索引即可实现，且完全适配「保留元素远少于无关元素」的场景，性能远高于逐行操作。 #

最高效方案：全零数组批量赋值 #

直接通过Numpy构造行、列索引数组，批量赋值需要保留的元素，时间复杂度仅和你需要保留的元素数量挂钩：

import numpy as np

# 构造行索引：每个行号重复i_array每行的元素个数次
rows = np.repeat(np.arange(z_array.shape[0]), i_array.shape[1])
# 构造列索引：直接将i_array展平即可
cols = i_array.ravel()

# 生成同维度全零数组
a_array = np.zeros_like(z_array)
# 批量赋值需要保留的元素
a_array[rows, cols] = z_array[rows, cols]

用你给出的示例代码测试，输出结果和预期的a_array完全一致。 #

关于numpy.ufunc.at的说明 #

np.ufunc.at确实可以实现该需求，但它的适用场景是索引存在重复、需要原地累加或重复赋值的情况。如果你的i_array每行不存在重复索引，普通高级索引赋值的性能比ufunc.at更高，无需额外使用。如果存在重复索引需要保留赋值结果，可以使用如下写法：

a_array = np.zeros_like(z_array)
# 两种写法均可，根据实际场景选择
np.put(a_array, np.ravel_multi_index((rows, cols), z_array.shape), z_array[rows, cols])
# 或ufunc.at写法
np.add.at(a_array, (rows, cols), z_array[rows, cols])

第三种实现思路：布尔掩码过滤 #

你可以生成布尔掩码数组，直接过滤原数组的元素，适合需要保留原z_array不被修改的场景：

mask = np.zeros(z_array.shape, dtype=bool)
mask[rows, cols] = True
a_array = z_array * mask

该方案的可读性很高，性能和第一种方案接近，缺点是会多一次数组乘法运算，超大规模数组下性能略低于第一种方案。 #

性能对比说明 #

你提到的第二种思路「直接在z_array基础上将无关位置置零」不推荐使用，因为无关元素数量远大于需要保留的元素，置零操作需要处理的元素量是第一种方案的几十倍，性能差距非常明显。

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