Possible to use np.where to check a condition in vector, but output rows in a 2D array

Viewed 39

I have a series and a dataframe. I want to check if the values in a series pass a condition, and modify the row of the dataframe if they do, otherwise leave as is.

NumPy has a broadcasting issue with this - is there another way to do this?

ser = pd.Series([74, 80, 24], pd.date_range(start='2020-01-01', periods=3, freq='D'))
test = pd.DataFrame([pd.Series([1, 2], index=['a', 'b'])] * len(ser), index=ser.index)
np.where(ser<50, (test*2), test)

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

I think a workaround would be to modify ser to be a dataframe with all equivalent columns, but it seems a little bit clunky.

1 Answers

Use broadcasting in NumPy, so they are not aligned by indices, only necessary same length of Series and DataFrame:

a = np.where(ser.to_numpy()[:, None]<50, (test*2), test)
print (a)
[[1 2]
 [1 2]
 [2 4]]
Related