python - 使用由整数 [0,...,L-1] 索引的额外列将 numpy 数组 (N,M,L) 转换为 (N*L,M+1)

Question

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

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 循环来做到这一点，但我相信还有更优雅的解决方案。

Answer 1

我首先对数组的轴重新排序，然后分配更大的结果数组，然后使用两个(广播)分配来设置新值:

import numpy as np

a = np.arange(6*3*5).reshape(6, 3, 5)  # shape (N, M, L)

aux = a.transpose(0, 2, 1)  # shape (N, L, M)
res = np.empty_like(a, shape=aux.shape[:-1] + (aux.shape[-1] + 1,))
res[..., :-1] = aux  # (N, L, M)-shaped slice
res[..., -1] = np.arange(aux.shape[1])  # (N, L)-shaped slice

# two different interpretations:
#res = res.reshape(-1, res.shape[-1])  # shape (N*L, M + 1)
res = res.transpose(0, 1, 2).reshape(-1, res.shape[-1])  # shape (N*L, M + 1)

在您的问题的两种解释中，后者(未注释版本)复制了您在评论中发布的“脏版本”。如果这确实是您所需要的，我们可以以放置 L 的方式进行原始转置。 -尺寸轴第一:

import numpy as np

a = np.arange(6*3*5).reshape(6, 3, 5)  # shape (N, M, L)

aux = a.transpose(2, 0, 1)  # shape (L, N, M)
res = np.empty_like(a, shape=aux.shape[:-1] + (aux.shape[-1] + 1,))
res[..., :-1] = aux  # (L, N, M)-shaped slice
res[..., -1] = np.arange(aux.shape[0])[:, None]  # (L, N)-shaped slice

res = res.reshape(-1, res.shape[-1])  # shape (N*L, M + 1)