在矩阵和对角矩阵之间进行矩阵乘法的更快方法？

Question

I am coding a computational package in python using numpy, in the package, I would do the matrix multiplication between an arbitrary large square matrix (eg of size 100*100) and a diagonal matrix of same size frequently.

I have an O(n^2) method, but I think that further improvement could be made.

"""
A is of size 100*100
B is a diagonal matrix 
want to do np.dot(A,B) quickly
"""
A=np.random.rand(100,100)
diag_elements=np.random.rand(100)
B=np.diag(diag_elements)

answer1= np.dot(A,B) ###O(n^3) method, quite slow

C=np.zeros((100,100)) 
C=C+diag_elements
answer2=np.multiply(A,C) ##O(n^2) method, 3times faster for n=100

The anwer2 is O(n^2) but I think it's not good enough, because the operation C+=diag_elements are wasting 1/3 time and could be avoided possibly.

I expect that some numpy function could do the matrix multiplication more elegently and faster. Could someone help me out?

Answer 1

为什么不简单地将A乘以对角线呢？

answer3 = np.multiply(A,diag_elements)