Haskell如何实现类似Python numpy的N维数组通用抽象索引功能？

核心问题本质 #

你遇到的类型不匹配问题根源在于：嵌套列表的每一层索引都会改变返回值的类型，而foldl要求折叠过程中累积值的类型固定，因此无法直接用foldl实现多层嵌套列表的通用索引。以下是几种可行的实现方案：

方案1：使用原生支持多维索引的数组库（推荐生产环境使用） #

Haskell生态里已经有成熟的任意维度数组库，最常用的是massiv，相比仅支持2维以下矩阵的hmatrix，它支持任意维度的张量操作，完全支持类似numpy的元组式索引，维度类型安全，还能自动优化内存排布，性能和numpy相当。
示例代码：

import Data.Massiv.Array as M

-- 定义一个3维Int数组，大小3x3x3
cube :: Array U Ix3 Int
cube = makeArrayR U Seq (Sz3 3 3 3) (\(i :. j :. k) -> i * 9 + j * 3 + k)

-- 直接用元组样式的索引取值，和numpy用法完全一致
main = print $ cube ! (0 :. 1 :. 2) -- 输出 0*9 +1*3 +2 = 5

它的索引类型IxN天然支持任意维度，编译期会检查索引维度和数组维度是否匹配，不会出现运行时维度错误。

方案2：手动实现嵌套列表的多态索引函数 #

如果你不想引入第三方依赖，只想对基础嵌套列表做通用索引，可以用Haskell的类型类实现多态索引，编译期自动匹配嵌套层级：

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverlappingInstances #-}

class Indexable a where
  (!) :: a -> [Int] -> a

instance {-# OVERLAPS #-} Indexable [a] where
  [] ! _ = error "索引越界"
  xs ! [] = xs
  xs ! (i:is) = (xs !! i) ! is

instance {-# OVERLAPS #-} Indexable a where
  x ! [] = x
  _ ! _ = error "索引长度超过数组维度"

-- 使用示例
main = do
  let cube = [[[1,2,3],[4,5,6],[7,8,9]], [[10,11,12],[13,14,15],[16,17,18]], [[19,20,21],[22,23,24],[25,26,27]]]
  print $ cube ! [0,1,2] -- 输出6

这个实现支持任意层级的嵌套列表，会自动根据索引长度逐层取值，如果索引长度和维度不匹配会在运行时报错。

方案3：类型安全的固定长度索引实现 #

如果你需要编译期保证索引长度和数组维度匹配，可以用GADT实现长度固定的索引列表，和维度对应的数组类型绑定，完全避免运行时维度错误：

{-# LANGUAGE GADTs #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE KindSignatures #-}

data Nat = Z | S Nat

data Vec :: Nat -> * -> * where
  Nil :: Vec Z a
  (:>) :: a -> Vec n a -> Vec (S n) a

infixr 5 :>

class NestedList (n :: Nat) a | n -> a where
  index :: a -> Vec n Int -> Element a

type family Element a where
  Element [a] = Element a
  Element a = a

instance NestedList Z a where
  index x Nil = x

instance NestedList n a => NestedList (S n) [a] where
  index xs (i :> is) = index (xs !! i) is

-- 使用示例
main = do
  let cube = [[[1,2,3],[4,5,6],[7,8,9]], [[10,11,12],[13,14,15],[16,17,18]], [[19,20,21],[22,23,24],[25,26,27]]]
  print $ index cube (0 :> 1 :> 2 :> Nil) -- 输出6

这种实现下如果索引长度和数组维度不匹配，编译阶段就会报错，不会带到运行时。

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