CCF-CSP真题202305-2《矩阵运算》（矩阵乘）

题目背景

是 Transformer 中注意力模块的核心算式，其中 Q、K 和 V 均是 n 行 d 列的矩阵， 表示矩阵 K 的转置，× 表示矩阵乘法。

问题描述

为了方便计算，顿顿同学将 Softmax 简化为了点乘一个大小为 n 的一维向量 W：

点乘即对应位相乘，记 为向量 W 的第 i 个元素，即将  第 i 行中的每个元素都与  相乘。

现给出矩阵 Q、K 和 V 和向量 W，试计算顿顿按简化的算式计算的结果。

输入格式

从标准输入读入数据。

输入的第一行包含空格分隔的两个正整数 n 和 d，表示矩阵的大小。

接下来依次输入矩阵 Q、K 和 W。每个矩阵输入 n 行，每行包含空格分隔的 d 个整数，其中第 i 行的第 j 个数对应矩阵的第 i 行、第 j 列。

最后一行输入 n 个整数，表示向量 W。

输出格式

输出到标准输出中。

输出共 n 行，每行包含空格分隔的 d 个整数，表示计算的结果。

样例输入

3 2
1 2
3 4
5 6
10 10
-20 -20
30 30
6 5
4 3
2 1
4 0 -5

样例输出

480 240
0 0
-2200 -1100

子任务

70 的测试数据满足：n≤100 且 d≤10；输入矩阵、向量中的元素均为整数，且绝对值均不超过 30。

全部的测试数据满足：n≤ 且 d≤20；输入矩阵、向量中的元素均为整数，且绝对值均不超过 1000。

提示

请谨慎评估矩阵乘法运算后的数值范围，并使用适当数据类型存储矩阵中的整数。

思路： 

1、看数据规格，全部的测试数据满足：n≤ 且 d≤20；输入矩阵、向量中的元素均为整数，且绝对值均不超过 1000。行数要远远超过列数，因此若采用原计算顺序会超时，只会得到70分，因此计算要适当有一点优化。

2、由于矩阵运算没有交换律但是有结合律，即(AB)C等于A(BC)，且Q*为n d * d n = n n，而*V为d n * n d = d d，数据规格中超过30%的测试案例n的大小要远远超过d，因此优化有效。

3、因此交换顺序后，后序的就是矩阵乘法的操作，也就是三层循环。

代码实现： 

#include<iostream>
#include<vector>

using namespace std;

int main() {
	vector<vector<int>>Q;
	vector<vector<int>>K;
	vector<vector<int>>V;
	vector<vector<long long>>res;
	vector<int>W;
	int n, d;
	cin >> n >> d;
	for (int i = 0; i < n; ++i) {          //输入矩阵Q
		vector<int>tmp;
		for (int j = 0; j < d; ++j) {
			int a;
			cin >> a;
			tmp.push_back(a);
		}
		Q.push_back(tmp);
	}
	for (int i = 0; i < n; ++i) {           //输入矩阵K
		vector<int>tmp;
		for (int j = 0; j < d; ++j) {
			int a;
			cin >> a;
			tmp.push_back(a);
		}
		K.push_back(tmp);
	}
	for (int i = 0; i < n; ++i) {           //输入矩阵V
		vector<int>tmp;
		for (int j = 0; j < d; ++j) {
			int a;
			cin >> a;
			tmp.push_back(a);
		}
		V.push_back(tmp);
	}
	for (int i = 0; i < n; ++i) {           //输入向量W
		int a;
		cin >> a;
		W.push_back(a);
	}
	vector<vector<long long>>KTV;
	for (int i = 0; i < d; ++i) {          //先计算K的转置乘以V，结果存放在矩阵KTV中 
		vector<long long>tmp;
		for (int j = 0; j < d; ++j) {
			long long sum = 0;
			for (int k = 0; k < n; ++k) {
				sum += K[k][i] * V[k][j];
			}
			tmp.push_back(sum);
		}
		KTV.push_back(tmp);
	}
	for (int i = 0; i < n; ++i) {           //再计算Q*KTV
		vector<long long>tmp;
		for (int j = 0; j < d; ++j) {
			long long sum = 0;
			for (int k = 0; k < d; ++k) {
				sum += Q[i][k] * KTV[k][j];
			}
			tmp.push_back(sum);
		}
		res.push_back(tmp);
	}
	for (int i = 0; i < n; ++i)              //最后计算向量W的每一个元素乘以矩阵中对应行的每一个元素
		for (int j = 0; j < d; ++j) {
			res[i][j] *= W[i];
		}
	for (int i = 0; i < res.size(); ++i) {   //输出结果矩阵
		for (int j = 0; j < res[0].size(); ++j) {
			cout << res[i][j] << " ";
		}
		cout << "\n";
	}
}