CSP2023_5 矩阵运算 Java

注意事项：因为n的规模远远小于d的，所以先算K的转置叉乘V，得到一个d×d的矩阵，时间复杂度O(d×n×d),时间和空间都会大大减小。

import java.util.Scanner;

public class Main {
    public static void main(String[] args) {
        Scanner in = new Scanner(System.in);
        int n = in.nextInt();
        int d = in.nextInt();//矩阵的大小
        int[][] Q = new int[n][d];
        int[][] K = new int[n][d];
        int[][] V = new int[n][d];//矩阵Q,K,V
        //输入矩阵
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < d; j++) {
                Q[i][j] = in.nextInt();
            }
        }
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < d; j++) {
                K[i][j] = in.nextInt();
            }
        }
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < d; j++) {
                V[i][j] = in.nextInt();
            }
        }
        //输入W
        int[] W = new int[n];
        for (int i = 0; i < n; i++) {
            W[i] = in.nextInt();
        }
        //转置矩阵K——》K1
        int[][] K1 = new int[d][n];
        for (int i = 0; i < d; i++) {
            for (int j = 0; j < n; j++) {
                K1[i][j] = K[j][i];
            }
        }
        //计算W(QXK1)V//矩阵乘法    先算K1XV
        long[][] temp = new long[d][d];
        for (int i = 0; i < d; i++) {
            for (int j = 0; j < d; j++) {
                for (int k = 0; k < n; k++) {
                    temp[i][j] += K1[i][k] * V[k][j];
                }
            }
        }

        //QX
        long[][] temp2 = new long[n][d];
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < d; j++) {
                for (int k = 0; k < d; k++) {
                    temp2[i][j]+=Q[i][k]* temp[k][j];
                }
            }
        }
        //W*
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < d; j++) {
                temp2[i][j]*=W[i];
            }
        }
        //打印
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < d; j++) {
                System.out.print(temp2[i][j]+" ");
            }
            System.out.println("");
        }

    }
}