填充numpy数组，使每一行与前一行偏移1

Question

我想弄清楚如何使用下面所示的模式填充数组:

     0000
0     000
00     00
000     0
0000      

例如:

[[1,2,3],
 [4,5,6],
 [7,8,9]]

将成为:

[[1,2,3,0,0],
 [0,4,5,6,0],
 [0,0,7,8,9]]

我知道如何使用手动循环，但我觉得可能有一个更快的方法，使用numpy的数组操作函数，我只是不知道如何做到这一点。

Answer 1

在Numpy中没有直接做这个的函数。您可以使用循环遍历行，但这种解决方案效率低下，除非2D数组非常大。

---

Numba解决方案一个非常有效的解决方案是使用Numba和琐碎循环:

import numba as nb
@nb.njit
def row_shifts_numba(arr):
    n, m = arr.shape
    out = np.zeros((n, n+m-1), dtype=arr.dtype)
    for i in range(n):
        for j in range(m):
            out[i, i+j] = arr[i, j]
    return out
data = np.array([[1,2,3],
                 [4,5,6],
                 [7,8,9]])
row_shifts_numba(data)

注意，由于插入式编译，第一次执行较慢。如果您不希望在第一次执行期间支付此编译时间，则可以在签名中指定数组的类型(例如。@nb.njit('int32[:,::1](int32[:,::1])')(其中int32为输入/输出数组类型，::1为轴连续)。

---

纯numpy替代方案

另一个解决方案是使用Numpy重塑技巧来生成输出的2D数组。这个想法是创建一个更大的二维数组，然后重塑它，以产生移位:

def row_shifts_numpy(arr):
    n, m = arr.shape
    out = np.zeros((n, n+m), dtype=arr.dtype)
    out[:n,:m] = arr
    return out.reshape(-1)[:-n].reshape(n, n+m-1)

由于(隐式)创建临时数组，这应该比Numba慢一点，但它是完全向量化的，并且只使用Numpy。

Answer 2

在创建具有所需形状的零数组之后，我们可以创建一个索引数组，其中的值逐行移动。然后用原始数组值填充零数组:

n, m = arr.shape
result = np.zeros((n, n+m-1), dtype=np.int64)
first_col_ind = np.array(np.arange(m))
ind = first_col_ind[:, None] + np.arange(n)
# [[0 1 2]
#  [1 2 3]
#  [2 3 4]]
result[np.arange(n), ind] = arr.T
# result:
# [[1 2 3 0 0]
#  [0 4 5 6 0]
#  [0 0 7 8 9]]