Picking rows from two NumPy arrays at random-CodePudding

Starting from two numpy arrays:

A = np.array([[1, 2], [2, 3], [3, 4], [4, 5]])
B = np.array([[9, 8], [8, 7], [7, 6], [6, 5]])

I would like to create a new array C picking, for each index, one row from the same index but randomly from A or B. The idea is that at each index of random_selector, if the value is higher than 0.1, then we chose the same-index row from A, otherwise, the same-index row from B.

random_selector = np.random.random(size=len(A))

C = np.where(random_selector > .1, A, B)

# example of desired result picking rows from respectively A, B, B, A:
# [[1, 2], [8, 7], [7, 6], [4, 5]]

Running the above code, however, produces the following error:

ValueError: operands could not be broadcast together with shapes (4,) (4,2) (4,2)

CodePudding user response：

Try adding a new dimension:

import numpy as np

A = np.array([[1, 2], [2, 3], [3, 4], [4, 5]])
B = np.array([[9, 8], [8, 7], [7, 6], [6, 5]])

random_selector = np.random.random(size=len(A))

C = np.where((random_selector > .1)[:, None], A, B)
print(C)

Output (of a single run)

[[1 2]
 [8 7]
 [3 4]
 [4 5]]