题解 | #实现Masked Multi-Head Self-Attention#

import math

n, d = map(int, input().split())

x = []

for _ in range(n):
    x.append(list(map(float, input().split())))

c = list(map(int, input().split()))

for ii in range(n):
    l = 0.0
    for jj in range(d):
        l += x[ii][jj] * x[ii][jj]
    l /= d
    low = math.sqrt(l)
    for jj in range(d):
        x[ii][jj] /= low

S = 0

def dot_mul(x1, x2):
    s = 0
    for iii in range(d):
        s += x1[iii] * x2[iii]
    return s


for jj in range(n):
    Aj = []
    for ii in range(jj):
        Aij = dot_mul(x[ii], x[jj]) / math.sqrt(d)
        Aj.append(Aij * Aij)
    Aj.sort(reverse=True)
    for ii in range(jj):
        S += Aj[ii]
        if ii == c[jj] - 1:
            break

print(round(100 * S))