python - Numpy 沿一个轴提取任意子数组

Question

例如，我有一个 3 维数组:

[[[ 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 30 31]
  [32 33 34 35]]]

和我想要的最终数组:

[[[ 0  1]
  [ 4  5]]

 [[18 19]
  [22 23]]

 [[26 27]
  [30 31]]]

有没有什么有效的方法可以在不使用 for 循环的情况下获取这样的数组？

问这个问题的原因是，如果我们想单独获取一个轴上的任意元素，例如:

[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]

我可以使用 a[np.arange(a.shape[0]), [2, 3, 1]] 获取数组 [ 2 7 9]所以我想知道当元素变成子数组时是否有类似的方式。

Answer 1

一个简单的想法是 a[[0,1,2],[0:2,1:3,0:2],[0:2,2:4,2:4]] ，但未实现。

解决方法可以通过 np.lib.stride_tricks.as_strided 获得.只需定义:

ast=np.lib.stride_tricks.as_strided(a,a.shape*2,a.strides*2)
#ast.shape is (3, 3, 4, 3, 3, 4).

然后你可以分别定义 block 的来源和大小:

In [4]: ast[[0,1,2],[0,1,0],[0,2,2],0,:2,:2]
Out[4]: 
array([[[ 0,  1],
        [ 4,  5]],

       [[18, 19],
        [22, 23]],

       [[26, 27],
        [30, 31]]])

一些解释:

• 来源:

您想找到以元素 0,18,26 开头的 block 。

可以通过以下方式找到它们在 reshape 数组中的索引:

In [316]: np.unravel_index([0,18,26],a.shape)
Out[316]: 
(array([0, 1, 2], dtype=int64),
 array([0, 1, 0], dtype=int64),
 array([0, 2, 2], dtype=int64))

ast[[0,1,2],[0,1,0],[0,2,2]] 是一个 (3,3,3,4) 数组。每个 (3,3,4) 数组都以一个选定的元素开头。

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,          30,          31],
         [         32,          33,          34,          35]]],


       [[[         18,          19,          20,          21],
         [         22,          23,          24,          25],
         [         26,          27,          28,          29]],

        [[         30,          31,          32,          33],
         [         34,          35,    23592960,       18335],
         [  697780028, -2147480064,   540876865,  1630433390]],

        [[ 2036429426,   538970664,   538976288,  1532698656],
         [  741355058,   808334368,   775168044,   874523696],
         [  744304686,   538976266,   538976288,   811278368]]],


       [[[         26,          27,          28,          29],
         [         30,          31,          32,          33],
         [         34,          35,    23592960,       18335]],

        [[  697780028, -2147480064,   540876865,  1630433390],
         [ 2036429426,   538970664,   538976288,  1532698656],
         [  741355058,   808334368,   775168044,   874523696]],

        [[  744304686,   538976266,   538976288,   811278368],
         [  539766830,   741355058,   808333600,   775036972],
         [  170679600,   538976288,   538976288,   774920992]]]])

如 documentation 中所述, as_strided 是一个危险的 hack，必须小心使用，因为如果使用不当，它可以访问不在数组中的元素。下一步将确保选择有效元素。

• 尺寸:

有趣的元素是每个第一 block 中左上角的四个元素。所以 ast[[0,1,2],[0,1,0],[0,2,2],0,:2,:2] 选择它们。

您还可以像这样定义大小的 block :bloc122=ast[...,0,:2,:2](bloc122.shape 是 (3, 3, 4, 2, 2) ) :

In [8]: bloc122[[0,1,2],[0,1,0],[0,2,2]]=0

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

       [[12, 13, 14, 15],
        [16, 17,  0,  0],
        [20, 21,  0,  0]],

       [[24, 25,  0,  0],
        [28, 29,  0,  0],
        [32, 33, 34, 35]]])