题解 | #注意力调度问题#

题目链接

题目描述

给定 $N$ 个 $D$ 维的特征向量 $X_0, ..., X_{N-1}$ ，以及每个位置 $j$ 的容量 $c_j$ 。你需要选择一系列“注意力”连边 $(i, j)$ （其中 $i < j$ ），使得每个位置 $j$ 接收的连边不超过 $c_j$ 条，并且所有选定连边的信息总量 $S$ 最大。

信息总量的计算步骤如下：

RMSNorm 归一化：对每个向量 $X_j$ 进行归一化得到 $\hat{X_j}$ 。
贡献度计算：每条连边 $(i, j)$ 的贡献度为 $a(i, j)$ ，总信息量 $S$ 是所有选定连边的贡献度平方之和，即 $S = \sum a(i, j)^2$ 。
目标：最大化 $S$ ，并输出 round(100 * S)。

解题思路

这是一个结合了向量运算和贪心策略的模拟题。题目的关键在于正确理解计算流程，并根据示例反推出一个与描述有出入的关键公式。

1. RMS 归一化

首先，按照题目定义对每个输入向量 $X_j$ 进行归一化。对于向量 $X_j = (x_1, ..., x_d)$ ：

计算各分量平方的均值：mean_sq = (x_1^2 + ... + x_d^2) / d
计算 rms = sqrt(mean_sq)
归一化向量 $\hat{X_j}$ 的每个分量为原分量除以 rms。

2. 修正贡献度公式

题目描述中，贡献度 $a(i, j)$ 被简单地描述为归一化后向量的点积。然而，如果我们直接使用 $a(i, j) = \text{dot}(\hat{X_i}, \hat{X_j})$ 来计算，其结果与示例给出的中间值和最终答案均不相符。

通过对示例数据进行反向推导，可以发现实际的贡献度公式为： $a(i, j) = \frac{\text{dot}(\hat{X_i}, \hat{X_j})}{\sqrt{d}}$ 即点积之后，还需要除以维度 $d$ 的平方根。后续的计算必须基于这个修正后的公式。

3. 贪心选择

题目明确指出：“对每个 $j$ 的选择彼此独立...贪心即最优”。这为我们提供了清晰的解题路径：

我们可以对每个接收位置 $j$ （从 $1$ 到 $N-1$ ）进行独立的决策。
对于一个固定的 $j$ ，我们计算所有可能的来源位置 $i$ （其中 $0 \le i < j$ ）到它的贡献度平方，即 $a(i, j)^2$ 。
将所有这些 $a(i, j)^2$ 的值存入一个列表。
对该列表进行降序排序。
从排序后的列表中，选取前 $c_j$ 个最大的值（如果列表长度小于 $c_j$ ，则全选）。
将这 $c_j$ 个值累加到总信息量 $S$ 中。

4. 汇总与输出

遍历完所有 $j$ 后，我们就得到了最大的总信息量 $S$ 。
最后，根据题目要求，计算 round(100 * S) 并输出。round 为标准四舍五入。

代码

cpp
java
python

#include <iostream>
#include <vector>
#include <cmath>
#include <numeric>
#include <algorithm>
#include <iomanip>

using namespace std;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(NULL);

    int n, d;
    cin >> n >> d;

    vector<vector<double>> vectors(n, vector<double>(d));
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < d; ++j) {
            cin >> vectors[i][j];
        }
    }

    vector<int> capacities(n);
    for (int i = 0; i < n; ++i) {
        cin >> capacities[i];
    }

    // 1. RMSNorm 归一化
    vector<vector<double>> normalized_vectors(n, vector<double>(d));
    for (int i = 0; i < n; ++i) {
        double sum_sq = 0.0;
        for (int j = 0; j < d; ++j) {
            sum_sq += vectors[i][j] * vectors[i][j];
        }
        double rms = sqrt(sum_sq / d);
        for (int j = 0; j < d; ++j) {
            normalized_vectors[i][j] = vectors[i][j] / rms;
        }
    }

    double total_s = 0.0;
    double sqrt_d = sqrt(d);

    // 2. 贪心选择
    for (int j = 1; j < n; ++j) {
        if (capacities[j] == 0) {
            continue;
        }

        vector<double> contributions_sq;
        for (int i = 0; i < j; ++i) {
            double dot_product = 0.0;
            for (int k = 0; k < d; ++k) {
                dot_product += normalized_vectors[i][k] * normalized_vectors[j][k];
            }
            double a_ij = dot_product / sqrt_d;
            contributions_sq.push_back(a_ij * a_ij);
        }

        sort(contributions_sq.rbegin(), contributions_sq.rend());

        for (int k = 0; k < min((int)contributions_sq.size(), capacities[j]); ++k) {
            total_s += contributions_sq[k];
        }
    }

    // 3. 输出
    cout << llround(100 * total_s) << "\n";

    return 0;
}

import java.util.Scanner;
import java.util.Vector;
import java.util.Collections;
import java.util.ArrayList;

public class Main {
    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        int n = sc.nextInt();
        int d = sc.nextInt();

        double[][] vectors = new double[n][d];
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < d; j++) {
                vectors[i][j] = sc.nextDouble();
            }
        }

        int[] capacities = new int[n];
        for (int i = 0; i < n; i++) {
            capacities[i] = sc.nextInt();
        }

        // 1. RMSNorm 归一化
        double[][] normalizedVectors = new double[n][d];
        for (int i = 0; i < n; i++) {
            double sumSq = 0.0;
            for (int j = 0; j < d; j++) {
                sumSq += vectors[i][j] * vectors[i][j];
            }
            double rms = Math.sqrt(sumSq / d);
            for (int j = 0; j < d; j++) {
                normalizedVectors[i][j] = vectors[i][j] / rms;
            }
        }

        double totalS = 0.0;
        double sqrtD = Math.sqrt(d);

        // 2. 贪心选择
        for (int j = 1; j < n; j++) {
            if (capacities[j] == 0) {
                continue;
            }

            ArrayList<Double> contributionsSq = new ArrayList<>();
            for (int i = 0; i < j; i++) {
                double dotProduct = 0.0;
                for (int k = 0; k < d; k++) {
                    dotProduct += normalizedVectors[i][k] * normalizedVectors[j][k];
                }
                double a_ij = dotProduct / sqrtD;
                contributionsSq.add(a_ij * a_ij);
            }

            contributionsSq.sort(Collections.reverseOrder());

            for (int k = 0; k < Math.min(contributionsSq.size(), capacities[j]); k++) {
                totalS += contributionsSq.get(k);
            }
        }

        // 3. 输出
        System.out.println(Math.round(100 * totalS));
    }
}

import math

def main():
    n, d = map(int, input().split())
    
    vectors = []
    for _ in range(n):
        vectors.append(list(map(float, input().split())))
        
    capacities = list(map(int, input().split()))

    # 1. RMSNorm 归一化
    normalized_vectors = []
    for vec in vectors:
        sum_sq = sum(x * x for x in vec)
        rms = math.sqrt(sum_sq / d)
        normalized_vectors.append([x / rms for x in vec])
        
    total_s = 0.0
    sqrt_d = math.sqrt(d)

    # 2. 贪心选择
    for j in range(1, n):
        if capacities[j] == 0:
            continue
            
        contributions_sq = []
        for i in range(j):
            dot_product = sum(normalized_vectors[i][k] * normalized_vectors[j][k] for k in range(d))
            a_ij = dot_product / sqrt_d
            contributions_sq.append(a_ij * a_ij)
            
        contributions_sq.sort(reverse=True)
        
        num_to_take = min(len(contributions_sq), capacities[j])
        total_s += sum(contributions_sq[:num_to_take])

    # 3. 输出
    # 标准四舍五入
    print(int(100 * total_s + 0.5))

if __name__ == "__main__":
    main()

算法及复杂度

算法：模拟、贪心
时间复杂度： $\mathcal{O}(N^2 \cdot D)$ 。主要瓶颈在于两层循环遍历所有可能的连边 $(i, j)$ ，并在内层计算 $D$ 维向量的点积。对于每个 $j$ ，还需要 $\mathcal{O}(N \log N)$ 的时间进行排序，但这部分通常小于点积计算的总耗时。
空间复杂度： $\mathcal{O}(N \cdot D)$ ，用于存储所有归一化后的向量。