从每个 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 loop

array([[[ 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 npa=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 npa = np.arange(2*3*5).reshape(2, 3, 5)np.array([a[0,:,2:4], a[1,:,1:3]])