I have predicted probabilities of test data for class 0 and class 1 in a 2D Numpy array. The true labels of the test data are in a 1D Numpy array. For all class 0 and class 1 records, I want to calculate the mean of class 0 and class 1 probabilities. I can use the "for" loop for this. Is there any Pythonic way for this calculation?

E.g.

psx =

array([[0.3, 0.7],

[0.6, 0.4],

[0.1, 0.9],

[0.8, 0.2]])

y_test=np.array([1,0,0,1])

The answer should be: array([0.35, 0.45])