Python中的矩阵对角线和leetcode

简介

在Python中，矩阵是一个常用的数据结构。由于矩阵中主对角线上的元素具有特殊意义，因此计算矩阵对角线和是一道常见的题目，而这个题目也出现在了LeetCode上。本文将介绍Python中如何计算矩阵对角线和，并提供LeetCode上的一道与之相关的题目。

矩阵对角线和的计算

在Python中，可以通过两种方式来计算矩阵对角线和。第一种方式是利用numpy库提供的函数。具体如下：

import numpy as np
matrix = np.array([[1,2,3],[4,5,6],[7,8,9]])
diagonal_sum = np.trace(matrix)
print(diagonal_sum)

通过调用numpy库的trace函数可以直接计算矩阵的对角线和。上述代码中，matrix是一个3*3的矩阵，diagonal_sum的值为15。

第二种方式是手动计算。具体如下：

matrix = [[1,2,3],[4,5,6],[7,8,9]]
diagonal_sum = sum([matrix[i][i] for i in range(len(matrix))])
print(diagonal_sum)

通过遍历矩阵的主对角线，将元素的值相加即可得到矩阵的对角线和。上述代码中，matrix是一个3*3的矩阵，diagonal_sum的值同样为15。

LeetCode题目

LeetCode上有一道与矩阵对角线和相关的题目，即题目编号为1572的“矩阵对角线元素的和”。题目描述如下：

给你一个正方形矩阵 mat，请你返回矩阵对角线元素的和。

我们定义矩阵的对角线为：mat[i][j] 其中 i == j。

示例 1：

输入：mat = [[1,2,3], [4,5,6], [7,8,9]] 输出：15

解释：对角线的和为：1 + 5 + 9 = 15。

示例 2：

输入：mat = [[1,1,1,1], [1,1,1,1], [1,1,1,1], [1,1,1,1]] 输出：4

解释：对角线的和为：1 + 1 + 1 + 1 = 4。

提示：

• n == mat.length == mat[i].length
• 1 <= n <= 100
• 1 <= mat[i][j] <= 100

解题思路如下：

class Solution:
    def diagonalSum(self, mat: List[List[int]]) -> int:
        n = len(mat)
        diagonal_sum = sum([mat[i][i] + mat[i][n-i-1] for i in range(n)])
        if n % 2 == 1:
            diagonal_sum -= mat[n//2][n//2]
        return diagonal_sum

其中，n表示矩阵的阶数，diagonal_sum表示对角线和。由于矩阵的对角线上会有某个元素被重复计算，因此需要特判当矩阵的阶数为奇数时，中心点不应重复计算两次。