Haskell中高效无坐标的8邻域计算实现需求

我们需要实现一个Haskell函数neighbors :: a -> [[a]] -> [[NonEmpty a]]，功能为：对矩形矩阵中的每个元素，生成包含自身（作为内部列表首元素）的8邻域列表，边界外的元素用指定默认值填充，每个邻域列表固定长度为9。

简洁实现方案 #

首先导入必要模块：

import Data.List (transpose)
import Data.NonEmpty (NonEmpty(..), toNonEmpty, toList)

优化后的augment填充函数 #

该函数负责给原矩阵的上下左右填充默认值，让原矩阵的每个元素都能处于3x3窗口的中心：

augment :: a -> [[a]] -> [[a]]
augment def mat = paddedRows
  where
    cols = length (head mat)
    emptyRow = replicate cols def
    paddedTop = emptyRow : mat
    paddedBottom = paddedTop ++ [emptyRow]
    paddedRows = map (\row -> def : row ++ [def]) paddedBottom

neighbors核心函数实现 #

利用滑动窗口和转置操作，高效生成每个元素的邻域：

neighbors :: a -> [[a]] -> [[NonEmpty a]]
neighbors def mat = concatMap extractWindows $ transpose $ map (slidingWindow 3) paddedMat
  where
    paddedMat = augment def mat
    -- 生成列表的所有长度为n的滑动窗口
    slidingWindow n xs = take (length xs - n + 1) $ map (take n) $ iterate tail xs
    -- 从3x3滑动窗口中提取并展平邻域列表
    extractWindows row = map (toNonEmpty . concat) $ slidingWindow 3 row

代码解释 #

1. 填充矩阵：augment先给原矩阵上下各添加一行默认值，再给每一行的首尾各添加一个默认值，得到一个比原矩阵大一圈的填充矩阵。
2. 滑动窗口生成：slidingWindow函数生成列表的所有连续n元素窗口，比如对[a,b,c,d]调用slidingWindow 3会得到[[a,b,c],[b,c,d]]。
3. 邻域提取：先对填充矩阵的每一行生成3元素窗口，转置后再生成3元素窗口，每个最终窗口对应原矩阵中一个元素的3x3邻域；将窗口展平后转成NonEmpty类型（首元素即为原矩阵的对应元素），最终拼接得到和原矩阵同结构的结果。

验证示例 #

运行以下代码：

main = do
  let mat = [[1,2,3],[4,5,6],[7,8,9]]
  mapM_ print mat
  putStrLn ""
  mapM_ (print . toList) $ neighbors 0 mat

输出与题目预期一致：

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

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

内容的提问来源于stack exchange，提问作者Brendan Langfield