基于另一数组实现Numpy元素迁移的高效方法求助

给定两个Numpy数组：

import numpy as np

data_arr = np.array([[0, 1, 2],
                     [1, 0, 0]])

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

需要按照以下规则迁移data_arr中的元素，生成同形状的全0数组new_data_arr：

• 若data_arr中元素为0，忽略迁移操作
• 若迁移目标位置在data_arr中非0，忽略迁移操作
• 若移动向量为(0,0)，元素保留在原位置
• 若目标位置超出数组范围，可选择忽略迁移或循环映射到数组另一侧
• 迁移时若目标位置在new_data_arr中已有值（来自之前的迁移），也忽略本次迁移

预期结果：

new_data_arr = [[1, 1, 0],
                [0, 0, 2]]

实际处理的数组规模极大，循环实现速度过慢，需要高效的矢量化实现。

现有循环实现代码：

def move(data_arr, move_arr):
    height, width = np.shape(data_arr)
    new_data_arr = np.zeros((height, width))
    for y in range(height):
        for x in range(width):
            target_y = y + move_arr[y, x][0]
            target_x = x + move_arr[y, x][1]
            if not np.all(move_arr[y, x] == 0) and data_arr[target_y, target_x] == 0 and new_data_arr[target_y, target_x] == 0:
                new_data_arr[target_y, target_x] = data_arr[y, x]
            elif data_arr[y, x] != 0:
                new_data_arr[y, x] = data_arr[y, x]
    return new_data_arr

# 测试
new_data_arr = move(data_arr, move_arr)
print(new_data_arr)

利用Numpy矢量化操作替代全数组嵌套循环，大幅提升处理速度，步骤如下：

1. 生成坐标网格与目标坐标 #

创建原位置坐标网格，结合move_arr计算目标坐标，同时处理边界循环映射：

height, width = data_arr.shape
# 生成原位置坐标网格
y, x = np.indices((height, width))
# 提取移动向量的y、x分量
dy, dx = move_arr[..., 0], move_arr[..., 1]
# 计算目标坐标，取模实现循环边界映射
target_y = (y + dy) % height
target_x = (x + dx) % width

2. 生成有效迁移掩码 #

根据规则筛选符合条件的迁移操作：

# 掩码：原位置元素非0
mask_data_nonzero = data_arr != 0
# 掩码：移动向量非(0,0)
mask_move_nonzero = (dy != 0) | (dx != 0)
# 掩码：目标位置在原data_arr中为0
mask_target_empty = data_arr[target_y, target_x] == 0
# 合并有效迁移掩码
mask_valid_move = mask_data_nonzero & mask_move_nonzero & mask_target_empty

3. 初始化并执行迁移 #

先保留无需迁移的元素，再按行优先顺序执行有效迁移（匹配原循环逻辑）：

new_data_arr = np.zeros_like(data_arr)
# 填充不需要迁移的元素（移动向量为0或目标位置不可迁移的非0元素）
mask_stay = mask_data_nonzero & ~mask_valid_move
new_data_arr[y[mask_stay], x[mask_stay]] = data_arr[mask_stay]

# 按行优先顺序处理有效迁移，确保目标位置未被填充
valid_indices = np.where(mask_valid_move)
sorted_indices = np.lexsort((valid_indices[1], valid_indices[0]))
sorted_y = valid_indices[0][sorted_indices]
sorted_x = valid_indices[1][sorted_indices]
sorted_target_y = target_y[sorted_y, sorted_x]
sorted_target_x = target_x[sorted_y, sorted_x]

# 仅对有效迁移元素循环，远快于全数组遍历
for sy, sx, ty, tx in zip(sorted_y, sorted_x, sorted_target_y, sorted_target_x):
    if new_data_arr[ty, tx] == 0:
        new_data_arr[ty, tx] = data_arr[sy, sx]

完整函数实现 #

import numpy as np

def vectorized_move(data_arr, move_arr, wrap_around=True):
    height, width = data_arr.shape
    new_data_arr = np.zeros_like(data_arr)
    
    # 生成坐标网格与移动分量
    y, x = np.indices((height, width))
    dy, dx = move_arr[..., 0], move_arr[..., 1]
    
    # 计算目标坐标，处理边界
    if wrap_around:
        target_y = (y + dy) % height
        target_x = (x + dx) % width
    else:
        target_y = y + dy
        target_x = x + dx
        # 过滤越界坐标
        mask_in_bounds = (target_y >= 0) & (target_y < height) & (target_x >=0) & (target_x < width)
    
    # 生成有效迁移掩码
    mask_data_nonzero = data_arr != 0
    mask_move_nonzero = (dy != 0) | (dx != 0)
    mask_target_empty = data_arr[target_y, target_x] == 0
    
    mask_valid_move = mask_data_nonzero & mask_move_nonzero & mask_target_empty
    if not wrap_around:
        mask_valid_move &= mask_in_bounds
    
    # 填充无需迁移的元素
    mask_stay = mask_data_nonzero & ~mask_valid_move
    new_data_arr[y[mask_stay], x[mask_stay]] = data_arr[mask_stay]
    
    # 按行优先执行有效迁移
    valid_indices = np.where(mask_valid_move)
    sorted_indices = np.lexsort((valid_indices[1], valid_indices[0]))
    sorted_y = valid_indices[0][sorted_indices]
    sorted_x = valid_indices[1][sorted_indices]
    sorted_target_y = target_y[sorted_y, sorted_x]
    sorted_target_x = target_x[sorted_y, sorted_x]
    
    for sy, sx, ty, tx in zip(sorted_y, sorted_x, sorted_target_y, sorted_target_x):
        if new_data_arr[ty, tx] == 0:
            new_data_arr[ty, tx] = data_arr[sy, sx]
    
    return new_data_arr

# 测试示例
data_arr = np.array([[0,1,2],[1,0,0]])
move_arr = np.array([[(0,0), (0,0), (1,0)], [(-1,0), (0,-1), (0,0)]])

new_data_arr = vectorized_move(data_arr, move_arr)
print(new_data_arr)
# 输出：
# [[1. 1. 0.]
#  [0. 0. 2.]]

性能说明 #

• 核心逻辑使用Numpy矢量化计算，避免了全数组嵌套循环
• 仅对有效迁移的少量元素执行循环，处理速度远优于原实现
• 支持边界循环映射或越界忽略两种模式
• 按行优先顺序处理迁移，与原循环逻辑完全一致，确保结果匹配

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