python - 矩阵中的对角线 - python

Question

我需要编写一个函数，该函数将返回在任何给定矩阵中形成对角线的数字之和。
作为一个 Python 新手，我遇到了一个问题。这是我的代码：

def diagonal(matrix):
    return sum([matrix[i][i] for i in range(len(matrix))])

我已经尝试了一段时间，但看不出有什么问题，因为它总是让我返回有关“列表索引超出范围”的错误报告。

我不允许导入numpy.

任何形式的帮助、提示、建议将不胜感激。

Answer 1

如果您确定您的矩阵是矩形的（len(matrix[i])对于 中的所有列表都相同matrix），那么只要您的维度较小，您就可以对列表求和：

def diagonal(matrix):
    return sum([matrix[i][i] for i in range(min(len(matrix[0]),len(matrix)))])

len(matrix)是矩阵的第一个维度，并且len(matrix[0])是第一行向量的维度，这是矩形矩阵的第二个维度。

Answer 2

当任一索引超出相应维度时，您必须停止，例如，您可以使用切片限制矩阵：

def diagonal_sum(matrix):
    row_size = len(matrix[0])
    return sum(row[i] for i, row in enumerate(matrix[:row_size]))

演示：

>>> diagonal_sum([[1,2],[3,4],[5,6]])
5

Answer 3

我认为diagonal不是为非方阵定义的。所以我们最好不要min只为了让代码返回一些东西而选择二维的。

那么如何取而代之：

def diagonal(matrix):
    try:
        return sum([matrix[i][i] for i in range(len(matrix))])
    except IndexError:
        print "Bad Matrix! :("