如何基于2D k索引边界数组高效填充3D numpy掩码数组？

问题背景 #

现有2D数组boundary_indices，其中非零值为k维度的边界索引。需要基于这些索引，为3D数组output的对应(i,j)位置填充两个区域：k索引上方（:k）赋值为7，下方（k:）赋值为8，仅处理boundary_indices中非零的(i,j)位置。原循环实现可得到正确结果，但数组规模增大时速度不足。

原初始化代码：

import numpy as np
boundary_indices = np.array([[0, 1, 2], [1, 2, 1], [0, 2, 0]])
ii, jj = np.where(boundary_indices > 0)  # 提取有效(i,j)索引
kk = boundary_indices[ii, jj]  # 对应位置的k边界值

原循环填充代码：

output = np.zeros((3, 3, 3), dtype=np.int64)
for i, j, k in zip(ii, jj, kk):
    output[i, j, :k] = 7
    output[i, j, k:] = 8

尝试直接使用花式索引output[ii, jj, :kk] = 7时，触发错误：TypeError: only integer scalar arrays can be converted to a scalar index，原因是numpy的花式索引不支持为每个(i,j)对指定不同长度的切片。

高效向量化实现 #

利用numpy的广播机制构造布尔掩码，实现批量赋值，避免循环：

import numpy as np

boundary_indices = np.array([[0, 1, 2], [1, 2, 1], [0, 2, 0]])
ii, jj = np.where(boundary_indices > 0)
kk = boundary_indices[ii, jj]

# 初始化3D输出数组
output = np.zeros((3, 3, 3), dtype=np.int64)
k_dim = output.shape[2]
# 生成k维度的所有索引，用于广播比较
k_indices = np.arange(k_dim)

# 构造上方区域掩码：每个(i,j)对应的k < kk[i]
mask_above = k_indices < kk[:, None]
# 为对应位置的上方区域赋值7
output[ii, jj, :] = np.where(mask_above, 7, output[ii, jj, :])

# 构造下方区域掩码：每个(i,j)对应的k >= kk[i]
mask_below = k_indices >= kk[:, None]
# 为对应位置的下方区域赋值8
output[ii, jj, :] = np.where(mask_below, 8, output[ii, jj, :])

原理说明 #

• k_indices < kk[:, None]通过广播将每个kk元素与k维度的所有索引比较，生成形状为(len(kk), k_dim)的布尔掩码，标记每个(i,j)位置需要赋值7的k索引。
• 利用np.where结合掩码完成批量赋值，所有操作均为numpy向量化运算，相比循环效率提升显著，尤其适合大规模数组。

验证结果与原循环实现完全一致：

print("output[:, :, 0]:")
print(output[:, :, 0])
print("\noutput[:, :, 1]:")
print(output[:, :, 1])
print("\noutput[:, :, 2]:")
print(output[:, :, 2])

输出：

output[:, :, 0]:
[[0 7 7]
 [7 7 7]
 [0 7 0]]

output[:, :, 1]:
[[0 8 7]
 [8 7 8]
 [0 7 0]]

output[:, :, 2]:
[[0 8 8]
 [8 8 8]
 [0 8 0]]

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