Return a numpy array, with numbers of elements specified in another array

Question

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 3x1 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?

swag2198 · Accepted Answer · 2020-11-22 18:44:20Z

1

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.

edited Nov 22, 2020 at 18:44

answered Nov 22, 2020 at 18:34

swag2198

2,7061 gold badge10 silver badges19 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Jingyang Wang Over a year ago

Thanks! This solution works beautifully when dimensions of A and B are large.

Collectives™ on Stack Overflow

Return a numpy array, with numbers of elements specified in another array

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related