从Haskell矩阵中提取对角线的最佳方法是什么？

Question

我被要求编写一个函数，该函数将提取存储为列表列表的矩阵的对角线。第一个版本是通过为列表建立索引来提取数字，但我很快得出结论，这对Haskell而言不是一个好算法，并编写了另一个函数：

getDiagonal :: (Num a) => [[a]] -> [a]
getDiagonal [[]]       = []
getDiagonal (xs:[])    = [head xs]
getDiagonal (x:xs)     = head x : getDiagonal (map tail xs)

由于我只是开始学习Haskell，所以我不确定它是不是以惯用的方式编写的，或者它是否会表现良好。

所以我的问题是，有没有更好的方法可以从存储在这种表示形式的矩阵中提取对角线，或者如果矩阵是使用更高阶的Haskell概念（如代数类型）表示的，那么是否可以构造出更好的算法？在像（（x：_）：xs）这样的模式匹配中解构列表还是与上面所示的head函数之间，在性能上有什么区别吗？

编辑：实际上，好奇的查询多于作业，他们在这里的技术大学不教函数编程（我认为这很可惜），但我会保留这个标签。

Answer 1

您可以将原始定义简化为：

mainDiagonal :: [[a]] -> [a]
mainDiagonal []     = []
mainDiagonal (x:xs) = head x : getDiagonal (map tail xs)

为此使用索引并没有多大错，这可以使您进一步简化为：

mainDiagonal xs = zipWith (!!) xs [0..]

基于数组的表示

您也可以使用由索引的Data.Array表示矩阵(i,j)。这使您几乎逐字使用主要对角线的数学定义：

import Data.Array

mainDiagonal :: (Ix i) => Array (i, i) e -> [e]
mainDiagonal xs = [ e | ((i,j),e) <- assocs xs, i == j ]

您可以这样使用：

-- n×n matrix helper
matrix n = listArray ((0,0),(n-1,n-1))

> mainDiagonal $ matrix 3 [1..]
[1,5,9]

效率

的先前定义mainDiagonal仍然无效：它仍然需要O（N²）个测试i == j。与zipWith版本类似，可以将其固定和概括如下：

mainDiagonal xs = (xs !) `map` zip [n..n'] [m..m']
                      where ((n,m),(n',m')) = bounds xs

此版本仅索引数组O（N）次。（此外，它还适用于矩形矩阵，并且独立于索引基数。）