单循环计算二维矩阵对角线和的效率及优化方案咨询

现有实现的效率评估 #

你的现有实现时间复杂度为O(n)（n为矩阵阶数），已经达到了该问题的理论最优时间复杂度——两条对角线合计共2n个元素（奇数阶矩阵中心元素重叠），必须至少遍历一次才能完成求和，不存在时间复杂度更低的解法。
但现有实现存在3个可优化的问题：

• 矩阵阶数硬编码为3，复用性极差，无法适配不同尺寸的方阵
• 奇数阶矩阵的中心元素会被重复累加，若你的业务场景要求中心元素仅计算一次，需要额外做判断修正
• 示例代码末尾输出的换行符存在语法错误，' '属于非法字符常量，正确写法为'\n'或std::endl

推荐的优化实现 #

1. 通用可复用版本（适配任意阶方阵，支持配置中心元素计算规则） #

#include <iostream>

// 参数说明：
// arr：输入方阵的二级指针
// n：方阵阶数
// allow_dup_center：奇数阶矩阵时，是否允许中心元素累加2次，默认允许
int calc_diag_sum(int** arr, int n, bool allow_dup_center = true) {
    int sum = 0;
    for (int i = 0; i < n; ++i) {
        sum += arr[i][i] + arr[n - 1 - i][i];
    }
    // 奇数阶且不允许重复累加中心元素时，减去多算的一次
    if (!allow_dup_center && n % 2 == 1) {
        sum -= arr[n / 2][n / 2];
    }
    return sum;
}

int main() {
    int arr[3][3] = {
        {1, 2, 3},
        {4, 5, 6},
        {7, 8, 9}
    };
    int* p_arr[3] = {arr[0], arr[1], arr[2]};
    // 输出30，和你示例的预期结果一致
    std::cout << calc_diag_sum(p_arr, 3) << '\n';
    // 输出25，中心元素5仅累加1次
    std::cout << calc_diag_sum(p_arr, 3, false) << '\n';
    return 0;
}

2. 固定尺寸小矩阵极致优化版本 #

如果你的使用场景固定为3x3矩阵，可以直接展开循环，甚至使用编译期计算，运行时完全无计算开销：

#include <iostream>

// 编译期计算3x3矩阵对角线和
constexpr int calc_3x3_diag_sum(const int arr[3][3]) {
    return arr[0][0] + arr[1][1] + arr[2][2] + arr[2][0] + arr[1][1] + arr[0][2];
}

int main() {
    constexpr int arr[3][3] = {
        {1, 2, 3},
        {4, 5, 6},
        {7, 8, 9}
    };
    // 编译阶段就已经算出sum=30，运行时直接输出常数
    constexpr int sum = calc_3x3_diag_sum(arr);
    std::cout << sum << '\n';
    return 0;
}

内容的提问来源于stack exchange，提问作者Itachi Uchiwa