如何利用3D索引数组为3D NumPy数组指定位置赋值？

问题描述 #

我有一个3D NumPy数组a，希望创建数组b，将a中绝对值最小和次小的位置设为0。我的实现代码如下：

import numpy as np

a = np.random.randn(100, 4, 4)
c = np.argsort(np.abs(a), axis=2)[..., :2]
b = np.ones(a.shape)

# I need to do b[c] = 0

但代码中的b[c] = 0无法正常运行。由于c仅包含0到3的索引值，b的第一维和第二维索引需要根据c中索引的位置自动推断，请问该如何解决这个问题？

解决方案 #

问题核心是c只提供了第三轴的索引，缺少前两轴的对应位置索引。需要生成前两轴的网格索引，与c组合后完成赋值，以下是两种高效实现方式：

方法一：使用np.indices生成完整索引网格 #

import numpy as np

a = np.random.randn(100, 4, 4)
c = np.argsort(np.abs(a), axis=2)[..., :2]
b = np.ones(a.shape)

# 生成第一、二维的索引网格
i, j = np.indices(a.shape[:2])
# 对指定位置赋值为0
b[i, j, c] = 0

• np.indices(a.shape[:2])生成两个形状为(100,4)的数组，分别对应第一维（100个样本）和第二维（4行）的索引，每个位置的i[x,y]就是第x个样本、第y行的索引值。
• 索引b[i, j, c]会通过广播机制，将c中每个位置的第三轴索引与前两轴的对应索引匹配，精准把每个样本每行中绝对值最小和次小的位置设为0。

方法二：使用np.ogrid生成内存更高效的开放网格 #

如果数组规模较大，推荐用np.ogrid生成开放网格，内存占用更低：

import numpy as np

a = np.random.randn(100, 4, 4)
c = np.argsort(np.abs(a), axis=2)[..., :2]
b = np.ones(a.shape)

# 生成开放索引网格
i, j = np.ogrid[:a.shape[0], :a.shape[1]]
# 完成赋值
b[i, j, c] = 0

np.ogrid生成的是一维数组的视图，而非完整的二维网格，在处理大数组时能显著节省内存。

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