Блог

Давайте рассмотрим этот простой пример:

 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-loops, но я уверен, что есть гораздо более элегантные решения.

Я бы сначала переупорядочил оси вашего массива, затем выделил больший результирующий массив, а затем использовал два (широковещательных) назначения для установки новых значений:

 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)
 

 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))
 

Сначала мы меняем местами последние 2 оси a с transpose , а затем reshape меняем их на желаемую форму, но с F порядком ортрана, чтобы соответствовать выходу. Дополнительная колонка получается с повторением 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]])