Python - 在每一行的不同位置切片数组

Question

我有一个 2D python 数组，我想以一种奇怪的方式切片 - 我想要一个恒定宽度的切片，从每一行的不同位置开始。如果可能的话，我想以矢量化的方式做到这一点。

例如，我有一个A=np.array([range(5), range(5)])看起来像的数组

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

我想按如下方式切片：每行 2 个元素，从位置 0 和 3 开始b=np.array([0,3])。起始位置存储在. 因此，所需的输出是：np.array([[0,1],[3,4]])即

array([[0, 1],
       [3, 4]])

我试图得到这个结果的显而易见的事情是，A[:,b:b+2]但这不起作用，而且我找不到任何可行的方法。

速度很重要，因为这将在循环中对较大的数组进行操作，而且我不想阻碍代码的其他部分。

Answer 1

方法#1：这是一种broadcasting获取所有索引然后用于advanced-indexing提取这些索引的方法 -

\n\n

def take_per_row(A, indx, num_elem=2):\n    all_indx = indx[:,None] + np.arange(num_elem)\n    return A[np.arange(all_indx.shape[0])[:,None], all_indx]\n

\n\n

样本运行 -

\n\n

In [340]: A\nOut[340]: \narray([[0, 5, 2, 6, 3, 7, 0, 0],\n       [3, 2, 3, 1, 3, 1, 3, 7],\n       [1, 7, 4, 0, 5, 1, 5, 4],\n       [0, 8, 8, 6, 8, 6, 3, 1],\n       [2, 5, 2, 5, 6, 7, 4, 3]])\n\nIn [341]: indx = np.array([0,3,1,5,2])\n\nIn [342]: take_per_row(A, indx)\nOut[342]: \narray([[0, 5],\n       [1, 3],\n       [7, 4],\n       [6, 3],\n       [2, 5]])\n

\n\n

方法#2：使用np.lib.stride_tricks.as_strided-

\n\n

from numpy.lib.stride_tricks import as_strided\n\ndef take_per_row_strided(A, indx, num_elem=2):\n    m,n = A.shape\n    A.shape = (-1)\n    s0 = A.strides[0]\n    l_indx = indx + n*np.arange(len(indx))\n    out = as_strided(A, (len(A)-num_elem+1, num_elem), (s0,s0))[l_indx]\n    A.shape = m,n\n    return out\n

\n\n

从矩阵中获取200每行的运行时测试2000x4000

\n\n

In [447]: A = np.random.randint(0,9,(2000,4000))\n\nIn [448]: indx = np.random.randint(0,4000-200,(2000))\n\nIn [449]: out1 = take_per_row(A, indx, 200)\n\nIn [450]: out2 = take_per_row_strided(A, indx, 200)\n\nIn [451]: np.allclose(out1, out2)\nOut[451]: True\n\nIn [452]: %timeit take_per_row(A, indx, 200)\n100 loops, best of 3: 2.14 ms per loop\n\nIn [453]: %timeit take_per_row_strided(A, indx, 200)\n1000 loops, best of 3: 435 \xc2\xb5s per loop\n

\n