Haskell高效实现Python列表指定位置赋值（埃氏筛场景）

首先得说，你之前尝试的两个Haskell实现性能拉胯是有原因的：

• 第一个用lens的elements (elem s)，每次判断元素是否在s里都要做线性遍历（因为s是列表的话elem是O(n)操作），对于筛法这种要处理大量数据的场景，复杂度直接爆炸。
• 第二个递归版本每次都要把s里的所有元素减1，还要拼接数组，不可变数组的每次修改都会生成新数组，这开销简直不忍直视。

针对埃氏筛的场景，我们需要原地修改或者高效批量更新的方案，下面给你几个最优解：

方案1：不可变数组的高效批量更新 #

如果坚持用不可变数组，推荐用Data.IntSet存储要标记的索引（查找是O(log n)），再配合数组的//操作批量更新：

import qualified Data.IntSet as IS
import Data.Array

setFalse :: IS.IntSet -> Array Int Bool -> Array Int Bool
setFalse indices arr = arr // [(k, False) | k <- IS.toList indices]

这个方案把多次单独更新合并成一次批量操作，比逐个修改高效很多，而且IntSet的查找/遍历性能比列表好太多。

方案2：用ST Monad的可变数组（性能最优） #

埃氏筛是典型的适合原地修改的场景，Haskell的ST monad允许我们安全地使用可变数组，性能和Python的原地修改几乎持平。

用STArray实现 #

import Control.Monad.ST
import Data.Array.ST

-- 原地将数组中指定索引的元素设为False
setFalseST :: [Int] -> STArray s Int Bool -> ST s ()
setFalseST indices arr = mapM_ (\k -> writeArray arr k False) indices

-- 基于STArray的埃氏筛示例
sieve :: Int -> [Int]
sieve n = runST $ do
  -- 初始化数组：2到n的所有数初始为True（认为是质数）
  arr <- newArray (2, n) True
  let markMultiples p = do
        -- 标记p的倍数（从p²开始，因为更小的倍数已经被之前的质数标记过）
        let multiples = [p*p, p*p+p .. n]
        setFalseST multiples arr
      loop p
        | p*p > n = do
            -- 提取所有标记为True的索引（质数）
            isPrime <- getElems arr
            return [i | (i, True) <- zip [2..n] isPrime]
        | otherwise = do
            prime <- readArray arr p
            if prime then markMultiples p else return ()
            loop (p+1)
  loop 2

用Unboxed Vector进一步提升性能 #

如果处理的n很大，推荐用Data.Vector.Unboxed.Mutable——它的内存布局更紧凑，缓存命中率更高，性能比STArray还要好：

import qualified Data.Vector.Unboxed.Mutable as VUM
import Control.Monad.ST

setFalseVec :: [Int] -> VUM.MVector s Bool -> ST s ()
-- 注意：向量索引从0开始，对应数字2，所以要把索引减2
setFalseVec indices vec = mapM_ (\k -> VUM.write vec (k-2) False) indices

sieveVec :: Int -> [Int]
sieveVec n = runST $ do
  -- 初始化向量：长度n-1，所有元素为True（对应数字2到n）
  vec <- VUM.replicate (n-1) True
  let markMultiples p = do
        let start = p*p
            multiples = [start, start+p .. n]
        setFalseVec multiples vec
      loop p
        | p*p > n = do
            vec' <- VUM.freeze vec
            return [i | (i, True) <- zip [2..n] (VUM.toList vec')]
        | otherwise = do
            let idx = p-2
            prime <- VUM.read vec idx
            if prime then markMultiples p else return ()
            loop (p+1)
  loop 2

为什么这些方案更快？ #

• 可变数组的writeArray/VUM.write是原地操作，没有不可变数组那样的复制开销。
• 避免了之前版本中反复遍历列表、修改列表元素的冗余操作。
• Unboxed Vector的内存布局更紧凑，能更好地利用CPU缓存，处理大规模数据时优势尤为明显。

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