题解 | 注意力调度问题

import sys
import numpy as np
data = list(map(int,sys.stdin.read().strip().split()))
n,d = data[:2]
X = np.array(list(map(float,data[2:2+n*d]))).reshape(n,d)
#print(X)
C = np.zeros(n)
C = np.array(list(map(int,data[-n:])))
#print(C)
import math

rms = np.sqrt((np.sum(X**2,axis=1))/d)
rms = rms[:,np.newaxis]
X = X / rms

A = np.zeros((n,n))
for j in range(n):
    for i in range(j):
        # 计算 a(i,j)   
        A[i,j] = X[i] @ X[j] / math.sqrt(d)
        A[i,j] = A[i,j] **2
# for  遍历所有的 r[i], sort ，选择c_j 个
S = 0
for j in range(n):

    contribution = []
    for i in range(j):
        contribution.append(A[i,j])
    contribution.sort(reverse=True)
    select_sum = sum(contribution[:min(C[j],len(contribution))])
    S += select_sum
S = round(S*100)
print(S)