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

Question

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

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

对我来说似乎有点 'ugly' 需要指定整个范围；此外，即使 arange 调用也需要时间：

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

有没有办法避免使用 arange？

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

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

我们可以滑动 windows 然后用起始索引沿着行和列索引那些以获得我们想要的输出。要获得这些 windows，我们可以利用 np.lib.stride_tricks.as_strided based scikit-image's view_as_windows. . This would be mostly inspired by .

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？

因为 windows 只是输入的视图，因此没有内存开销。只是在最后一步，在获取输出时，我们需要额外的内存 space 来保存所需的切片，无论如何都是必需的。

样本运行-

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

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

How is one supposed to use the arange function when indexing an ndarray?

python

numpy

numpy-slicing

numpy-ndarray