Suppose I have two numpy arrays A
, B
, with A.shape = (2,4,3)
and B.shape = (2,4)
:
A = [[[-1 -1 0] [-1 0 0] [-1 0 0] [ 0 -1 0]] [[-1 0 0] [-1 0 0] [-1 1 0] [ 0 0 0]]] B = [[0 2 0 3] [1 0 0 0]]
Now I would like to get a new array C
with C.shape = (2+3+1,3) = (6,3)
, such that each integer in B
specifies the number of the corresponding 3×1 array in A
. In other words, A[i,j,:]
should appear B[i,j]
times in C
. In our case,
C = [[-1 0 0] [-1 0 0] [ 0 -1 0] [ 0 -1 0] [ 0 -1 0] [-1 0 0]]
What is the fastest way to obtain C
?
Advertisement
Answer
I can think of doing this by reshaping A
into a simple 2D array of which the rows A[i]
, are to be repeated according to the count in B[i]
.
Note that B
is also flattened to be a 1D array.
import numpy as np A = np.array([[[-1, -1, 0], [-1, 0, 0], [-1, 0, 0], [ 0, -1, 0]], [[-1, 0, 0], [-1, 0, 0], [-1, 1, 0], [ 0, 0, 0]]]) B = np.array([[0, 2, 0, 3], [1, 0, 0, 0]]) A = A.reshape(-1, A.shape[2]) # A.shape becomes (8, 3) B = B.reshape(-1) # B.shape becomes (8, ) C = np.repeat(A, B, axis = 0) >>>C
Output:
array([[-1, 0, 0], [-1, 0, 0], [ 0, -1, 0], [ 0, -1, 0], [ 0, -1, 0], [-1, 0, 0]])
Here is how np.repeat()
works, in case you need reference.