独立移动 numpy 数组的行

Question

这是这里提出的问题的扩展（引用如下）

我有一个矩阵（准确地说是二维 numpy ndarray）：

A = np.array([[4, 0, 0],
              [1, 2, 3],
              [0, 0, 5]])

我想根据另一个数组中的滚动值独立滚动 A 的每一行：

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

也就是说，我想这样做：

print np.array([np.roll(row, x) for row,x in zip(A, r)])

[[0 0 4]
 [1 2 3]
 [0 5 0]]

有没有办法有效地做到这一点？也许使用花哨的索引技巧？

接受的解决方案是：

rows, column_indices = np.ogrid[:A.shape[0], :A.shape[1]]

# Use always a negative shift, so that column_indices are valid.
# (could also use module operation)
r[r < 0] += A.shape[1]
column_indices = column_indices - r[:,np.newaxis]

result = A[rows, column_indices]

我基本上想做同样的事情，除非当索引滚动“超过”行的末尾时，我希望行的另一侧用 NaN 填充，而不是将值移动到“前面” "以周期性方式排列。

也许以np.pad某种方式使用？但我不知道如何让它以不同的数量填充不同的行。

Answer 1

受独立于矩阵的滚动行解决方案的启发，这是一个基于np.lib.stride_tricks.as_strided-

from skimage.util.shape import view_as_windows as viewW

def strided_indexing_roll(a, r):
    # Concatenate with sliced to cover all rolls
    p = np.full((a.shape[0],a.shape[1]-1),np.nan)
    a_ext = np.concatenate((p,a,p),axis=1)

    # Get sliding windows; use advanced-indexing to select appropriate ones
    n = a.shape[1]
    return viewW(a_ext,(1,n))[np.arange(len(r)), -r + (n-1),0]

样品运行 -

In [76]: a
Out[76]: 
array([[4, 0, 0],
       [1, 2, 3],
       [0, 0, 5]])

In [77]: r
Out[77]: array([ 2,  0, -1])

In [78]: strided_indexing_roll(a, r)
Out[78]: 
array([[nan, nan,  4.],
       [ 1.,  2.,  3.],
       [ 0.,  5., nan]])