How to select rows from two different Numpy arrays conditionally?-CodePudding

I have two Numpy 2D arrays and I want to get a single 2D array by selecting rows from the original two arrays. The selection is done conditionally. Here is the simple Python way,

import numpy as np

a = np.array([4, 0, 1, 2, 4])
b = np.array([0, 4, 3, 2, 0])
y = np.array([[0, 0, 0, 0],
              [0, 0, 0, 1],
              [0, 0, 1, 0],
              [0, 0, 1, 1],
              [0, 0, 1, 0]])
x = np.array([[0, 0, 0, 0],
              [1, 1, 1, 0],
              [1, 1, 0, 0],
              [1, 1, 1, 1],
              [0, 0, 1, 0]])
z = np.empty(shape=x.shape, dtype=x.dtype)
for i in range(x.shape[0]):
    z[i] = y[i] if a[i] >= b[i] else x[i]

print(z)

Looking at numpy.select, I tried, np.select([a >= b, a < b], [y, x], -1) but got ValueError: shape mismatch: objects cannot be broadcast to a single shape. Mismatch is between arg 0 with shape (5,) and arg 1 with shape (5, 4).

Could someone help me write this in a more efficient Numpy manner?

CodePudding user response：

This should do the trick, but it would be helpful if you could show an example of your expected output:

>>> np.where((a >= b)[:, None], y, x)
array([[0, 0, 0, 0],
       [1, 1, 1, 0],
       [1, 1, 0, 0],
       [0, 0, 1, 1],
       [0, 0, 1, 0]])