如何在numpy中优雅地通过子数组重新排列数组？

Posted 2023-02-14

【中文标题】如何在numpy中优雅地通过子数组重新排列数组？

【英文标题】：How to rearrange an array by subarray elegantly in numpy?

【发布时间】：2017-06-13 18:17:38

【问题描述】：

假设我有一个 3-D 数组：

[[[0,1,2],
  [0,1,2],
  [0,1,2]],

 [[3,4,5],
  [3,4,5],
  [3,4,5]]]

我想按列重新排列：

[[0,1,2,3,4,5],
 [0,1,2,3,4,5],
 [0,1,2,3,4,5]]

对于本质上具有任意形状和深度的 3-D np.array 来说，一个优雅的 python numpy 代码会是什么？ 是否有一种绕过for循环的快速方法？我所做的所有方法都是非常临时和粗暴的，它们基本上太慢而且没用......

谢谢！！

【问题讨论】：

我想有人会说我正试图通过列索引来解开这个 3D 数组...

【参考方案1】：

交换轴并重塑 -

a.swapaxes(0,1).reshape(a.shape[1],-1)

示例运行 -

In [115]: a
Out[115]: 
array([[[0, 1, 2],
        [0, 1, 2],
        [0, 1, 2]],

       [[3, 4, 5],
        [3, 4, 5],
        [3, 4, 5]]])

In [116]: a.swapaxes(0,1).reshape(a.shape[1],-1)
Out[116]: 
array([[0, 1, 2, 3, 4, 5],
       [0, 1, 2, 3, 4, 5],
       [0, 1, 2, 3, 4, 5]])

【参考方案2】：

使用 einops：

einops.rearrange(a, 'x y z -> y (x z) ')

我建议根据上下文（例如时间、高度等）为轴（而不是 x y z）赋予有意义的名称。这样可以很容易理解代码的作用

In : einops.rearrange(a, 'x y z -> y (x z) ')
Out:
array([[0, 1, 2, 3, 4, 5],
       [0, 1, 2, 3, 4, 5],
       [0, 1, 2, 3, 4, 5]])