索引ndarray时应该如何使用arange函数？

Question

假设我想为每一行从不同的列中选择一个值。然后，我可能会这样做：

a = np.arange(12).reshape(3, 4)
columns = np.array([1, 2, 0])
a[np.arange(a.shape[0]), columns]

在我看来，需要指定整个范围有点“难看”；此外，即使是arange 调用也需要时间：

%timeit np.arange(int(1e6))
1.03 ms ± 15.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

有没有办法避免使用 arange？

概括上述问题；如何不选择单个值，而是为每一行选择不同的相邻列集（每组大小相同）？我想避免创建许多手动范围，如下所示：

rows = np.array([0, 2])
start_values = np.array([0, 1])
window_length = 3
column_ranges = np.array(list(map(lambda j: np.arange(j, j + window_length), start_values)))

现在，我看到使用上述列范围的唯一方法是像这样索引：

a[rows, :][:, column_ranges][np.arange(len(rows)), np.arange(len(rows)), :]

理想情况下，我想使用a[:, columns] 代替a[np.arange(a.shape[0]), columns] 和a[:, columns:columns + window_length] 代替a[rows, :][:, column_ranges][np.arange(len(rows)), np.arange(len(rows)), :] 之类的符号。

Answer 1

我们可以得到滑动窗口，然后用开始索引沿着行和列索引那些我们想要的输出。要获得这些窗口，我们可以利用基于scikit-image's view_as_windows 的np.lib.stride_tricks.as_strided。 More info on use of as_strided based view_as_windows。这主要是受到this post 的启发。

from skimage.util.shape import view_as_windows

def windows_per_row_vas(arr, rows, cols, W):
    w = view_as_windows(a,(1,W))[...,0,:]
    return w[rows,cols]

如果你想通过使用 np.lib.stride_tricks.as_strided 的粗略实现来弄脏你的手 -

def windows_per_row_strided(arr, rows, cols, W):
    strided = np.lib.stride_tricks.as_strided 
    m,n = arr.shape
    s0,s1 = arr.strides
    windows = strided(arr, shape=(m,n-W+1,W), strides=(s0,s1,s1))
    return windows[rows, cols]

为什么使用views/strided？

因为窗口只是输入视图，因此没有内存开销。只是在最后一步，当得到输出时，我们需要额外的内存空间来保存所需的切片，无论如何都是需要的。

示例运行 -

In [9]: a
Out[9]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])

In [10]: rows = np.array([0, 2])
    ...: start_values = np.array([0, 1])
    ...: window_length = 3

In [11]: windows_per_row_strided(a, rows, start_values, window_length)
Out[11]: 
array([[ 0,  1,  2],
       [ 9, 10, 11]])


In [29]: windows_per_row_vas(a, rows, start_values, window_length)
Out[29]: 
array([[ 0,  1,  2],
       [ 9, 10, 11]])