Как использовать np.cumsum для ускорения расчета средней точности?

Question

У меня есть,

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]])

Я хочу рассчитать на карте K, которую я написал как go, как показано ниже,

        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

Это в основном дает формулу расчета, как показано ниже : (1 + 0 + 2/3) / 3 + (0 + 0 + 1/3) / 2

Есть ли какие-либо предположения, что я мог бы использовать np.cumsum, чтобы ускорить мой al go? Я предполагаю, что он был оптимизирован, и я не вижу здесь никакой комнаты для улучшения?

Заранее спасибо.

Answer 1

надеюсь, это поможет вам (я стараюсь избегать циклов):

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