如何使用Python numpy对此表达式进行矢量化？

Question

给定numpy数组P(N),D(N,N),我想计算数组A(N):

A_i = SUM(P_i - P_j) where j is such that D[i][j] is True
A_i = 0 otherwise

所以,

A_i = N*P_i - (P_m + P_n+...) where D[i][m], D[i][n],... are all True

我可以使用双循环和enumerate()编写上面的内容,但它很慢.只是想知道如何用numpy ufuncs来写这个.

Answer 1

如果我理解正确的问题,这是一种方法.

A = P*D.sum(axis=1)-D.dot(P)