NumPy 1.10能否向量化含索引变量的函数？附三维数组需求

当然可以实现向量化处理！完全不用嵌套循环就能搞定这个需求，而且在NumPy 1.10里也能完美运行。核心思路是利用NumPy的索引和广播机制，一次性处理所有符合条件的位置，效率比循环高得多。

下面是具体的实现步骤和代码示例：

核心思路 #

1. 先定位数组A中所有为True的元素的三维索引；
2. 定义需要偏移的方向（比如三维空间中的6个邻域方向，或者你只需要的+1方向）；
3. 通过广播机制生成所有邻域位置的索引；
4. 过滤掉超出数组B边界的无效索引；
5. 一次性将B中有效邻域位置设为True。

代码实现 #

import numpy as np

# 生成示例数据（你可以替换成自己的A和B数组）
A = np.random.rand(5, 5, 5) > 0.7  # 随机三维布尔数组，模拟你的A
B = np.zeros((7, 7, 7), dtype=bool)  # 更大的目标数组B

# 步骤1：获取A中所有True元素的三维索引
i, j, k = np.where(A)

# 步骤2：定义偏移方向——这里是6个邻域方向（±1在每个轴）
# 如果只需要每个维度的+1步，改成offsets = np.array([[1,0,0], [0,1,0], [0,0,1]])即可
offsets = np.array([[1,0,0], [-1,0,0], [0,1,0], [0,-1,0], [0,0,1], [0,0,-1]])

# 步骤3：生成所有邻域的坐标
# 将A的True索引组合成(N, 3)的坐标数组
coords = np.stack([i, j, k], axis=1)
# 通过广播给每个坐标加上所有偏移，得到所有邻域坐标
neighbor_coords = coords[:, None, :] + offsets[None, :, :]
# 把多维数组展平成(N*6, 3)的二维数组，方便后续处理
neighbor_coords = neighbor_coords.reshape(-1, 3)

# 步骤4：过滤掉超出B边界的无效坐标
b_shape = B.shape
# 检查每个维度的索引是否在合法范围内
valid_mask = (neighbor_coords[:, 0] >= 0) & (neighbor_coords[:, 0] < b_shape[0]) & \
             (neighbor_coords[:, 1] >= 0) & (neighbor_coords[:, 1] < b_shape[1]) & \
             (neighbor_coords[:, 2] >= 0) & (neighbor_coords[:, 2] < b_shape[2])
# 筛选出有效的坐标
valid_coords = neighbor_coords[valid_mask]

# 步骤5：将B中对应位置设为True
B[valid_coords[:, 0], valid_coords[:, 1], valid_coords[:, 2]] = True

关键说明 #

• 整个过程完全没有嵌套循环，全部依赖NumPy的向量化操作，在NumPy 1.10中完全支持（np.stack、广播、索引操作都是该版本已有的功能）；
• 如果你的需求只是每个维度的**+1步**（而非正负两个方向），只需修改offsets数组，保留三个正方向的偏移即可；
• 代码中加入了边界过滤，避免出现索引越界的问题，确保所有操作都在B的范围内进行。

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