Numpy：根据另一个数组移动值

Question

我已经尝试搜索几个小时了，但找不到解决方案，所以我希望有人可以帮助我。

假设我有两个如下所示的 numpy 数组：

data_arr = [[0 1 2]
            [1 0 0]]

move_arr = [[(0,0)  (0,0)  (1,0)]
            [(-1,0) (0,-1) (0,0)]]

我需要找到一种方法来根据

data_arr

中的值移动

move_arr

中的值。其规则如下：如果

data_arr

中的值为 0，或者移动的目的地不为 0，则忽略该移动。

预期结果是：

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

请注意，实际上，我正在使用的数组非常庞大，而且这是性能关键的代码。因此，我不能只使用 for 循环。

所以我不得不问，这个，或者至少类似的东西，是否可以很快实现。如果是的话，那又如何呢？

我已经尝试过使用 for 循环，虽然它确实有效，但速度太慢了。

Answer 1

假设 numpy 数组作为输入，您可以使用

numba

有效地通过循环执行此操作：

from numba import jit

@jit(nopython=True)
def move(data_arr, move_arr):
    out = data_arr.copy()
    for i in range(data_arr.shape[0]):
        for j in range(data_arr.shape[1]):
            i2, j2 = move_arr[i, j]
            i2 += i
            j2 += j
            if out[i, j] != 0 and out[i2, j2] == 0:
                out[i2, j2] = data_arr[i, j]
                out[i, j] = data_arr[i2, j2]
    return out

out = move(data_arr, move_arr)

输出：

array([[1, 1, 0],
       [0, 0, 2]])

使用的输入：

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)]])