CCF CSP题解：矩阵运算（202305-2）

链接和思路

OJ链接：传送门

本题要求计算1个公式：
$$\left(\mathbf{W}\cdot (\mathbf{Q}\times {\mathbf{K}}^T)\right)\times \mathbf{V}$$
其中，$\mathbf{Q}$、$\mathbf{K}$和$\mathbf{V}$均是$n$行$d$列的矩阵，${\mathbf{K}}^T$，表示矩阵$\mathbf{K}$的转置， $\times$表示矩阵乘法。$\cdot$为点乘，即对应位相乘，记${\mathbf{W}}^{(i)}$为向量$\mathbf{W}$的第$i$个元素，即将$(\mathbf{Q}\times {\mathbf{K}}^T)$第$i$行中的每个元素都与${\mathbf{W}}^{(i)}$相乘。

本题有2点需要注意，否则只能过70%的样例：

• 使用int会导致溢出，可使用long long表示数据。
• 如果按照公式给出的顺序计算，复杂度为$O(d{n}^2)$，注意到$n$远大于$d$，因此应该修改运算顺序，优化到$O(d^{2}n)$。

由于注意到矩阵乘法${\mathbf{A}}_{n\times m}\times {\mathbf{B}}_{m\times k}$的复杂度是$O(nmk)$，因此我们尽可能要让$m$更小，于是原式的计算顺序可以改变为：
$$\left(\mathbf{W}\cdot (\mathbf{Q}\times {\mathbf{K}}^T)\right)\times \mathbf{V}=\mathbf{W}\cdot \left(\mathbf{Q}\times ({\mathbf{K}}^T\times \mathbf{V})\right)$$
调整矩阵乘法顺序在矩阵乘法计算优化中是十分常见的，如果是一连串任意给定的矩阵相乘，可以用动态规划得到最优的矩阵运算次序。此外，使用行优先的方式比列优先更能充分利用缓存命中率，但是由于已经满分，因此在本题中我们没有在这个方向继续优化。

AC代码

#include <iostream>
#include <vector>

using namespace std;

void print_vector(const vector<vector<long long>> &arr) {
    
    
    for (int i = 0; i < arr.size(); i++) {
    
    
        for (int j = 0; j < arr[0].size(); j++) {
    
    
            if (j != 0)
                cout << " ";
            cout << arr[i][j];
        }
        cout << endl;
    }
}

int main() {
    
    
    int n, d;
    cin >> n >> d;
    vector<vector<long long>> q(n), k(n), v(n);
    vector<long long> w(n);
    for (int i = 0; i < n; ++i) {
    
    
        q[i].resize(d);
        for (int j = 0; j < d; ++j) {
    
    
            cin >> q[i][j];
        }
    }

    for (int i = 0; i < n; ++i) {
    
    
        k[i].resize(d);
        for (int j = 0; j < d; ++j) {
    
    
            cin >> k[i][j];
        }
    }
    for (int i = 0; i < n; ++i) {
    
    
        v[i].resize(d);
        for (int j = 0; j < d; ++j) {
    
    
            cin >> v[i][j];
        }
    }

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

    //kv: d x d
    vector<vector<long long>> kv(d);
    for (int i = 0; i < d; ++i) {
    
    
        kv[i].resize(d);
    }
    for (int i = 0; i < d; ++i) {
    
    
        for (int j = 0; j < d; ++j) {
    
    
            for (int l = 0; l < n; ++l) {
    
    
                kv[i][j] += k[l][i] * v[l][j];
            }
        }
    }

    //qkv: n x d
    for (int i = 0; i < n; ++i) {
    
    
        for (int j = 0; j < d; ++j) {
    
    
            k[i][j] = 0;
            for (int l = 0; l < d; ++l) {
    
    
                k[i][j] += q[i][l] * kv[l][j];
//                printf("k[%d][%d]=%d\n", i, j, k[i][j]);
            }
        }
    }

    // wqkv: n x d
    for (int i = 0; i < n; i++)
        for (int j = 0; j < d; ++j)
            k[i][j] *= w[i];
    print_vector(k);

    return 0;
}