矩阵中的对角线 - 蟒蛇

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! :("