题解 | 矩阵乘法计算量估算

括号内字母可大于2

import java.util.*;

public class Main {
    public static void main(String[] args) {

        // 输入
        Scanner in = new Scanner(System.in);
        int n = in.nextInt(); // 矩阵个数

        // 矩阵数组
        Matrix[] matrices = new Matrix[n];
        for (int i = 0; i < n; i++) {
            int row = in.nextInt(); // 矩阵行数
            int col = in.nextInt(); // 矩阵列数
            matrices[i] = new Matrix(row, col);
        }

        String formula = in.next(); // 读取运算式（next()跳过换行符）

        int crossNum = 0; // 总运算量
        Stack<Character> stack = new Stack<>(); // 用于处理括号和矩阵

        for (int i = 0; i < formula.length(); i++) {
            char ch = formula.charAt(i);

            //输入为'('
            if (ch == '(') {
                stack.push(ch); // 压入左括号
            } 
            
            //输入为字母（矩阵）
            else if (ch >= 'A' && ch <= 'Z') {

                // 如果栈顶是字母（矩阵），则计算乘法，后将结果矩阵压回栈
                if (!stack.isEmpty() && stack.peek() != '(') {
                    char top = stack.pop();
                    crossNum += matrices[top - 'A'].cross(matrices[ch - 'A']);
                    stack.push(top); 
                } 
                // 如果栈顶不是字母（矩阵），则直接压入矩阵
                else {
                    stack.push(ch); 
                }
            } 
            
            //输入为')'，弹出栈顶的矩阵
            else {
   
                char top = stack.pop();
                stack.pop(); // 弹出对应的左括号

                // 如果栈顶是字母（矩阵），则计算乘法，后将结果矩阵压回栈
                if (!stack.isEmpty() && stack.peek() != '(') {
                    char prev = stack.pop();
                    crossNum += matrices[prev - 'A'].cross(matrices[top - 'A']);
                    stack.push(prev); 
                } 
                //如果栈顶不是字母（矩阵），则直接压入矩阵
                else {
                    stack.push(top);
                }
            }
        }

        System.out.println(crossNum); // 输出总运算量
    }

    //矩阵类
    static class Matrix {
        int row;
        int col;

        public Matrix(int row, int col) {
            this.row = row;
            this.col = col;
        }

        // 计算两个矩阵相乘的运算量，并更新当前矩阵的列数
        public int cross(Matrix other) {
            int result = row * col * other.col;
            col = other.col; // 更新当前矩阵的列数
            return result;
        }
    }
}