从每个 numpy 行中选择不同的切片

Question

我有一个 3d 张量，我想从 dim=2 中选择不同的切片。类似的东西a[[0, 1], :, [slice(2, 4), slice(1, 3)]]。a=np.arange(2*3*5).reshape(2, 3, 5)array([[[ 0,  1,  2,  3,  4],        [ 5,  6,  7,  8,  9],        [10, 11, 12, 13, 14]],       [[15, 16, 17, 18, 19],        [20, 21, 22, 23, 24],        [25, 26, 27, 28, 29]]])# then I want something like a[[0, 1], :, [slice(2, 4), slice(1, 3)]]# that gives me np.stack([a[0, :, 2:4], a[1, :, 1:3]]) without a for looparray([[[ 2,  3],        [ 7,  8],        [12, 13]],       [[16, 17],        [21, 22],        [26, 27]]])我已经看到了，这不是我想要的。

Answer 1

您可以advanced indexing按照此处的说明使用。您将必须传递[0, 1]您的情况下的行 ID 以及列 ID2, 3和1, 2. 这里的2,3意思[2:4]和1, 2意思[1:3]

import numpy as np

a=np.arange(2*3*5).reshape(2, 3, 5)

rows = np.array([[0], [1]], dtype=np.intp)

cols = np.array([[2, 3], [1, 2]], dtype=np.intp)

aa = np.stack(a[rows, :, cols]).swapaxes(1, 2)

# array([[[ 2,  3],

#         [ 7,  8],

#         [12, 13]],

#        [[16, 17],

#         [21, 22],

#         [26, 27]]])

另一种避免swapaxes并以所需格式获得结果的等效方法是

aa = np.stack(a[rows, :, cols], axis=2).T

我想出的第三种方法是传递索引列表。这里[0, 0]将对应[2,3]和[1, 1]将对应[1, 2]。这swapaxes只是为了获得您想要的输出格式

a[[[0,0], [1,1]], :, [[2,3], [1,2]]].swapaxes(1,2)

Answer 2

一个解法 ...

import numpy as np

a = np.arange(2*3*5).reshape(2, 3, 5)

np.array([a[0,:,2:4], a[1,:,1:3]])