首页 > 解决方案 > 将 numpy 数组 (N,M,L) 转换为 (N*L,M+1) ,其中额外的列由整数 [0,...,L-1] 索引

问题描述

让我们考虑这个简单的例子:

import numpy as np

a=np.arange(90)
a=a.reshape(6,3,5)

我想得到一个b形状数组(6 * 5,3 + 1 = 4)

b[0:6,0]=a[:,0,0]
b[0:6,1]=a[:,1,0]
b[0:6,2]=a[:,2,0]
b[0:6,3]=0

b[6:12,0]=a[:,0,1]
b[6:12,1]=a[:,1,1]
b[6:12,2]=a[:,2,1]
b[6:12,3]=1
...

我可以使用 for 循环来做到这一点,但我确信还有更优雅的解决方案。

标签: pythonarraysnumpy

解决方案


a_new = a.transpose(0, 2, 1).reshape(N*L, M, order="F")
extra_column = np.repeat(np.arange(L), N)
b = np.column_stack((a_new, extra_column))

a我们首先交换with的最后 2 个轴,transpose然后reshape将其交换为所需的形状,但使用Fortran 顺序以匹配输出。额外的列产生重复np.arange(L)并添加column_stack


样品运行:

>>> N, M, L = 6, 3 ,5
>>> a = np.arange(N*M*L).reshape(N, M, L)
>>> # above operations...
>>> b

array([[ 0,  5, 10,  0],
       [15, 20, 25,  0],
       [30, 35, 40,  0],
       [45, 50, 55,  0],
       [60, 65, 70,  0],
       [75, 80, 85,  0],
       [ 1,  6, 11,  1],
       [16, 21, 26,  1],
       [31, 36, 41,  1],
       [46, 51, 56,  1],
       [61, 66, 71,  1],
       [76, 81, 86,  1],
       [ 2,  7, 12,  2],
       [17, 22, 27,  2],
       [32, 37, 42,  2],
       [47, 52, 57,  2],
       [62, 67, 72,  2],
       [77, 82, 87,  2],
       [ 3,  8, 13,  3],
       [18, 23, 28,  3],
       [33, 38, 43,  3],
       [48, 53, 58,  3],
       [63, 68, 73,  3],
       [78, 83, 88,  3],
       [ 4,  9, 14,  4],
       [19, 24, 29,  4],
       [34, 39, 44,  4],
       [49, 54, 59,  4],
       [64, 69, 74,  4],
       [79, 84, 89,  4]])

推荐阅读