注意力调度问题

题意

给定长度为 $n$ 的序列，每个位置 $j$ 有一个 $d$ 维特征向量 $X_j$ 和一个容量 $c_j$ 。先对每个向量做 RMSNorm 归一化，然后对任意 $(i, j)$ （ $i < j$ ）计算贡献 $\frac{(\hat{x}_i \cdot \hat{x}_j)^2}{d}$ 。每个位置 $j$ 最多选 $c_j$ 条来自前序位置的入边，要求选边使得总贡献 $S$ 最大，输出 $\text{round}(100 \times S)$ 。

思路

拆开来看，每个位置 $j$ 的选边是独立的——位置 $j$ 选谁做入边，不影响位置 $k$ （ $k \neq j$ ）的可选集合。所以不存在全局竞争，每个位置各自贪心就行。

具体做法：

对每个向量做 RMSNorm： $\text{rms}_j = \sqrt{\frac{\sum_{k} x_{j,k}^2}{d}}$ ，归一化向量 $\hat{x}_j = \frac{X_j}{\text{rms}_j}$ 。
对每个位置 $j$ ，枚举所有 $i < j$ ，算出贡献 $\frac{(\hat{x}_i \cdot \hat{x}_j)^2}{d}$ 。
把这些贡献从大到小排序，取前 $c_j$ 个累加到 $S$ 。
最后输出 $\text{round}(100 \times S)$ 。

为什么每个位置独立贪心就是最优？

关键在于：入边选择没有"排他性"。位置 $i$ 被位置 $j$ 选为入边，不妨碍位置 $k$ （ $k > i$ ）也选 $i$ 做入边。容量限制 $c_j$ 只约束" $j$ 自己最多收几条边"，不约束" $i$ 最多被引用几次"。所以每个 $j$ 的决策互不干扰，各自选最大的 $c_j$ 个贡献即可。

拿样例验证： $n=4, d=2$ ，归一化后 $\hat{x}_0=[1,1]$ 、 $\hat{x}_1=[\sqrt{2},0]$ 、 $\hat{x}_2=[0,\sqrt{2}]$ 、 $\hat{x}_3=[1,1]$ 。

各对贡献（除以 $d=2$ ）：

$(i,j)$	$\hat{x}_i \cdot \hat{x}_j$	$(\cdot)^2 / d$
$(0,1)$	$\sqrt{2}$	$1$
$(0,2)$	$\sqrt{2}$	$1$
$(0,3)$	$2$	$2$
$(1,2)$	$0$	$0$
$(1,3)$	$\sqrt{2}$	$1$
$(2,3)$	$\sqrt{2}$	$1$

容量 $c = [0, 1, 1, 2]$ ：

$j=0$ ： $c_0 = 0$ ，不选。
$j=1$ ： $c_1 = 1$ ，选最大的 $(0,1) = 1$ 。
$j=2$ ： $c_2 = 1$ ，选最大的 $(0,2) = 1$ 。
$j=3$ ： $c_3 = 2$ ，选最大的两个 $(0,3) = 2$ 和 $(1,3) = 1$ （或 $(2,3) = 1$ ，一样大）。

$S = 1 + 1 + 2 + 1 = 5$ ，输出 $500$ 。

复杂度

时间： $O(n^2 \cdot d + n^2 \log n)$ ，对每个 $j$ 枚举 $i$ 算点积 $O(nd)$ ，排序 $O(n \log n)$
空间： $O(nd)$

代码

Python3
C++

import math

def solve():
    n, d = map(int, input().split())
    vecs = []
    for _ in range(n):
        x = list(map(float, input().split()))
        vecs.append(x)
    caps = list(map(int, input().split()))

    normed = []
    for x in vecs:
        rms = math.sqrt(sum(v * v for v in x) / d)
        normed.append([v / rms for v in x])

    S = 0.0
    for j in range(n):
        if caps[j] == 0:
            continue
        scores = []
        for i in range(j):
            dot = sum(normed[i][k] * normed[j][k] for k in range(d))
            scores.append(dot * dot / d)
        scores.sort(reverse=True)
        take = min(caps[j], len(scores))
        S += sum(scores[:take])

    print(round(100 * S))

solve()

#include <bits/stdc++.h>
using namespace std;

int main(){
    int n, d;
    scanf("%d%d", &n, &d);
    vector<vector<double>> v(n, vector<double>(d));
    for(int i = 0; i < n; i++){
        double ss = 0;
        for(int j = 0; j < d; j++){
            scanf("%lf", &v[i][j]);
            ss += v[i][j] * v[i][j];
        }
        double rms = sqrt(ss / d);
        for(int j = 0; j < d; j++) v[i][j] /= rms;
    }
    vector<int> c(n);
    for(int i = 0; i < n; i++) scanf("%d", &c[i]);

    double S = 0;
    for(int j = 1; j < n; j++){
        if(c[j] == 0) continue;
        vector<double> scores;
        for(int i = 0; i < j; i++){
            double dot = 0;
            for(int k = 0; k < d; k++) dot += v[i][k] * v[j][k];
            scores.push_back(dot * dot / d);
        }
        sort(scores.rbegin(), scores.rend());
        int take = min(c[j], (int)scores.size());
        for(int t = 0; t < take; t++) S += scores[t];
    }
    printf("%lld\n", llround(100.0 * S));
    return 0;
}