如何使用 np.cumsum 加快平均精度计算？

Question

I have,

scores = np.array([[0.9, 0.8, 0.6, 0.5, 0.4], [0.5, 0.4, 0.31, 0.21, 0.4 ]])
labels = np.array([[1, 0, 1, 1, 0], [0, 0, 0, 1, 1]])

I want to calculate at K map, which I wrote an algo as below,

        k=3
        mean_ap = 0
        n = len(scores)
        for i in range(n):
            cum = ap = 0.0
            idx = np.argsort(-scores[i])
            used_label = labels[i][idx][:k]
            m = sum(labels[i])
            for j, label in enumerate(used_label):
                cum += label
                ap += cum * label / (j + 1)
            mean_ap += ap / min(m, k)
        val = mean_ap / n

It basically gives calculation formula like below: (1 + 0 + 2 / 3) / 3 + ( 0 + 0 + 1 / 3) / 2

Any suggestion that I could use np.cumsum to speed up my algo? I assume it has been optimized and I don't see any enhancement room here?

Thanks in advance.

Answer 1

hope this can help you (I try to avoid for loops):

k = 3
n = len(scores)
m = labels.sum(axis=1)

idx = np.argsort(-scores)
used_label = labels[:,idx][np.arange(0,n),np.arange(0,n),:k]

val = (np.cumsum(used_label, axis=1)*used_label / 
       np.arange(1,k+1) / 
       np.min([m,np.repeat(k,n)],axis=0).reshape(-1,1)).sum(axis=1).sum() / n