Transform numpy array (N,M,L) into (N*L,M+1) with the extra column indexed by integer [0,...,L-1]

user_11635654 0
Viewed 59

Let us consider this simple example:

import numpy as np

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

I would like to get an array b of shape (6*5,3+1=4) with

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

I can do it with for-loops but I'm sure there are much more elegant solutions.

numpy arrays python
2 Answers

I'd first reorder the axes of your array, then allocate the larger resulting array, then use two (broadcasting) assignments to set the new values:

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)

Of the two interpretations of your question, the latter (uncommented version) reproduces your "dirty version" posted in a comment. If this is really what you need, we could do the original transpose in a way that puts the L-sized axis first:

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

We first swap the last 2 axes of a with transpose and then reshape it to desired shape but with Fortran order to match the output. Extra column is produced with repeated np.arange(L) and added with column_stack.

Sample run:

>>> 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]])
user_93321870
Related