如何用向量化方法移动3D NumPy数组中的非零值？

问题背景 #

我们有一个大规模稀疏NumPy数组（实际为1000×1000×1000网格，仅含10~1000个非零值1），需要根据对应的位移数组（3D场景下为shift_x、shift_y、shift_z），将每个1移动到目标位置，且支持循环环绕边界。以下用2D场景举例说明需求：

示例输入 #

原始位置数组：

import numpy as np

positions = np.array([[0, 0, 0, 0, 0],
                      [0, 0, 1, 0, 0],
                      [0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 1],
                      [0, 0, 0, 0, 0]])

位移规则数组：

shift_x = np.array([[0, 0,  0, 0, 0],
                    [0, 0, -1, 0, 0],
                    [0, 0,  0, 0, 0],
                    [0, 0,  0, 0, 1],
                    [0, 0,  0, 0, 0]])

shift_y = np.array([[0, 0, 0, 0,  0],
                    [0, 0, 1, 0,  0],
                    [0, 0, 0, 0,  0],
                    [0, 0, 0, 0, -1],
                    [0, 0, 0, 0,  0]])

预期输出 #

位移后的新位置数组：

new_positions = np.array([[0, 1, 0, 0, 0],
                          [0, 0, 0, 0, 0],
                          [0, 0, 0, 0, 0],
                          [0, 0, 0, 0, 0],
                          [1, 0, 0, 0, 0]])

循环逐个移动的方法在3D大规模场景下效率极低，需要基于NumPy的向量化操作实现高效处理。

向量化解决方案 #

核心思路是提取所有非零值的坐标，结合位移数组计算目标坐标，再通过坐标赋值生成新数组，全程避免循环：

步骤1：提取原始非零坐标 #

利用np.where快速定位所有1的位置：

# 2D场景：获取y, x坐标（NumPy数组索引为行优先，即y在前，x在后）
y_orig, x_orig = np.where(positions == 1)

# 3D场景对应：
# z_orig, y_orig, x_orig = np.where(positions == 1)

步骤2：获取对应位移量 #

根据原始坐标提取位移数组中的对应值：

# 2D场景
dx = shift_x[y_orig, x_orig]
dy = shift_y[y_orig, x_orig]

# 3D场景对应：
# dx = shift_x[z_orig, y_orig, x_orig]
# dy = shift_y[z_orig, y_orig, x_orig]
# dz = shift_z[z_orig, y_orig, x_orig]

步骤3：计算目标坐标（处理循环环绕） #

利用模运算实现边界循环：

# 2D场景
x_new = (x_orig + dx) % positions.shape[1]
y_new = (y_orig + dy) % positions.shape[0]

# 3D场景对应：
# x_new = (x_orig + dx) % positions.shape[2]
# y_new = (y_orig + dy) % positions.shape[1]
# z_new = (z_orig + dz) % positions.shape[0]

步骤4：生成新位置数组 #

初始化全零数组，通过目标坐标赋值：

# 2D场景
new_positions = np.zeros_like(positions)
new_positions[y_new, x_new] = 1

# 3D场景对应：
# new_positions = np.zeros_like(positions)
# new_positions[z_new, y_new, x_new] = 1

完整2D示例代码 #

import numpy as np

# 初始化输入数组
positions = np.array([[0, 0, 0, 0, 0],
                      [0, 0, 1, 0, 0],
                      [0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 1],
                      [0, 0, 0, 0, 0]])

shift_x = np.array([[0, 0,  0, 0, 0],
                    [0, 0, -1, 0, 0],
                    [0, 0,  0, 0, 0],
                    [0, 0,  0, 0, 1],
                    [0, 0,  0, 0, 0]])

shift_y = np.array([[0, 0, 0, 0,  0],
                    [0, 0, 1, 0,  0],
                    [0, 0, 0, 0,  0],
                    [0, 0, 0, 0, -1],
                    [0, 0, 0, 0,  0]])

# 向量化处理流程
y_orig, x_orig = np.where(positions == 1)
dx = shift_x[y_orig, x_orig]
dy = shift_y[y_orig, x_orig]

x_new = (x_orig + dx) % positions.shape[1]
y_new = (y_orig + dy) % positions.shape[0]

new_positions = np.zeros_like(positions)
new_positions[y_new, x_new] = 1

print(new_positions)

方案优势 #

• 效率极高：针对稀疏数组仅处理非零值，3D场景下即使1000个非零值，也只需执行几次数组操作，远快于循环遍历全量元素。
• 天然适配3D：只需扩展坐标维度和位移数组的提取逻辑，无需修改核心流程。
• 边界处理简洁：模运算直接实现循环环绕，无需额外判断边界条件。

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