Replace column with another array using python-CodePudding

I have one A array of (5x5) elements and a second B array of (5x1). Now I want to replace the column of A with B. My array is

A=array([[1, 2, -3, 4, 5],[1, 2, -3, -4, 5],[-1, 2, -3, -4, 5],[-1, -2, -3, -4, 5],[-1, -2, -3, -4, -5]])
and 
B=array([-10. , -11.5, -13. , -14.5, -16. ])

So it should replace every column with B when the first negative number encounters. For example, in the first column of A, the first negative element is at A[2][0] (-1), so the first column of A replaces with B[2], in the second column the first negative element is at A[3][1] and that column replace with B[3], etc. So my final output is

A=(alpha[2],alpha[3],alpha[0],alpha[1],alpha[4])= (-13.0, -14.5, -10.0, -11.5, -16.0)

So far my code is

A=array([[1, 2, -3, 4, 5],[1, 2, -3, -4, 5],[-1, 2, -3, -4, 5],[-1, -2, -3, -4, 5],[-1, -2, -3, -4, -5]])
B=np.linspace(-10,-16,5)
for i in range(len(B)):
  A[:,i][A[:,i]<0]=alpha[i]

But this is not giving the correct answer.

Thank you in advance!

CodePudding user response：

You could use argmax to find the first row for each column in which there is a negative value, and then use that result as an index into B:

import numpy as np

A = np.array([[1, 2, -3, 4, 5],[1, 2, -3, -4, 5],[-1, 2, -3, -4, 5],[-1, -2, -3, -4, 5],[-1, -2, -3, -4, -5]])
B = np.linspace(-10, -16, 5)

A2 = B[np.argmax(A < 0, axis=0)]

Output:

array([-13. , -14.5, -10. , -11.5, -16. ])