时间戳

首页 > 解决方案 > 我可以编写更有效的代码来在 Python 中获取对角线总和吗？

python - 我可以编写更有效的代码来在 Python 中获取对角线总和吗？

问题描述

我已经通过使用 for loops of row then column 解决了这个问题。但我想知道是否有更多 Pythonic 方式？谢谢。

问题 - 给出一个二维矩阵，得到两个对角线和。例如：

matrix = [[1,2,3],
          [4,5,6],
          [7,8,9]]
get_diag(matrix) # 25

标签： pythonmatrix

解决方案

您可以在展平矩阵上使用切片来对迭代器执行求和（即不使用索引）：

from itertools import chain,islice
def sumDiag(M):
    D  = len(M)
    S  = sum(islice(chain.from_iterable(M),0,None,D+1))           # main
    S += sum(islice(chain.from_iterable(reversed(M)),0,None,D+1)) # inv
    return S - M[D//2][D//2]*(D%2)  # subtract center if double counted

输出：

matrix = [[1,2,3],
          [4,5,6],
          [7,8,9]]

print(sumDiag(matrix)) #25


matrix = [[1,2,3,3],
          [4,5,6,0],
          [7,8,9,2],
          [4,0,2,2]]

print(sumDiag(matrix)) #38

或者，您可以生成索引并使用它们来获取要添加的项目。不确定这是否更“Pythonic”：

def sumDiag(M):
    return sum(matrix[r][r]+matrix[-r-1][r] for r in range(len(M))) \
           - M[len(M)//2][len(M)//2]*(len(M)%2)

推荐阅读