Haskell与Python堆（优先队列）性能对比及Haskell慢因排查

我用Python和Haskell分别实现了基于堆（优先队列）的问题解决方案，代码如下。测试发现当数据量达到10^5级别时，Haskell代码运行速度显著变慢，请求解释原因并给出优化方案。

Python实现代码 #

import heapq

_, volume, *ppl = map(int, open("input.txt").read().split())
ppl = [-x for x in ppl]
heapq.heapify(ppl)
for i in range(volume):
    current = (- heapq.heappop(ppl)) // 10
    heapq.heappush(ppl, -current)

print(sum(-x for x in ppl), file=open("output.txt", 'w'))

Haskell实现代码 #

import Data.List hiding (insert, singleton)

type Rank = Int

data Heap a = Tip | Node Rank Int (Heap Int) (Heap Int)

rank Tip = 0
rank (Node r _ _ _) = r

fromList :: [Int] -> Heap Int
fromList [] = Tip
fromList (x:xs) = foldl' (flip insert) (singleton x) xs

makeHeap :: Int -> Heap Int -> Heap Int -> Heap Int
makeHeap x a b = if rank a >= rank b then Node (rank b + 1) x a b
                                     else Node (rank a + 1) x b a

empty :: Heap a
empty = Tip

singleton :: Int -> Heap Int
singleton x = Node 1 x Tip Tip

insert :: Int -> Heap Int -> Heap Int
insert x = merge (singleton x)

merge :: Heap Int -> Heap Int -> Heap Int
merge l Tip = l
merge Tip r = r
merge h1@(Node _ x l1 r1) h2@(Node _ y l2 r2) =
  if x > y then makeHeap x l1 (merge r1 h2)
            else makeHeap y l2 (merge h1 r2)

-- | O(1).
peek :: Heap Int -> Int
peek Tip = 0
peek (Node _ x _ _) = x

-- | O(1), but evaluating the second element of the tuple has same complexity
-- of `merge`.
extract :: Heap Int -> (Int, Heap Int)
extract (Node _ x a b) = (x, merge a b)

toSum :: Heap Int -> Int
toSum Tip            = 0
toSum (Node _ x a b) = x + toSum a + toSum b 

solve :: Int -> Heap Int -> Int
solve 0 heap = toSum heap
solve _ (Node _ 0 _ _) = 0
solve sips ppl = solve (sips - 1) (insert (val `div` 10) h)
    where (val, h) = extract ppl
    

main :: IO()
main = do
    line <- readFile "input.txt"
    let _:volume:ppl = map(read :: String -> Int) $ words line
    let heap = fromList ppl
    writeFile "output.txt" $ show $ solve volume heap

测试结果片段 #

1. 内置库vs自定义实现：Python的heapq是C实现的高效优先队列，heapify操作是线性时间O(n)；而你的Haskell代码是手动实现的左偏堆，fromList通过逐个insert构建堆，时间复杂度为O(n log n)，纯Haskell实现的常数开销远高于C实现。
2. 惰性求值的开销：Haskell默认的惰性求值会在递归的merge和solve过程中产生大量未求值的thunk（延迟计算的表达式），数据量较大时，这种内存和计算开销会被显著放大。
3. 数据结构严格性不足：自定义Heap的Node字段是惰性的，每次访问子堆或值都可能触发额外计算，没有利用严格求值避免不必要的延迟。
4. 缺乏编译优化：如果编译时未开启-O2等优化选项，GHC不会进行严格化、内联、冗余计算消除等优化，导致代码运行效率低下。

1. 使用高效内置优先队列库 #

改用Data.PQueue.Max（来自priority-queue包），这是经过优化的最大堆实现，性能远超手动实现：

import qualified Data.PQueue.Max as MaxPQ

solve :: Int -> MaxPQ.MaxQueue Int -> Int
solve 0 pq = MaxPQ.foldl' (+) 0 pq
solve _ pq | MaxPQ.findMax pq == 0 = 0
solve sips pq = solve (sips - 1) (MaxPQ.insert (val `div` 10) pq')
    where (val, pq') = MaxPQ.deleteFindMax pq

main :: IO()
main = do
    line <- readFile "input.txt"
    let _:volume:ppl = map(read :: String -> Int) $ words line
    let pq = MaxPQ.fromList ppl
    writeFile "output.txt" $ show $ solve volume pq

2. 优化自定义堆的构建效率 #

将fromList改为线性时间的堆构建，通过分组合并子堆实现：

fromList :: [Int] -> Heap Int
fromList [] = Tip
fromList xs = mergePairs $ map singleton xs
    where
        mergePairs [] = Tip
        mergePairs [h] = h
        mergePairs (h1:h2:hs) = merge (merge h1 h2) (mergePairs hs)

此修改将fromList的时间复杂度从O(n log n)降至O(n)，大幅提升堆构建速度。

3. 开启编译优化 #

编译时添加-O2参数，让GHC进行深度优化：

ghc -O2 your_code.hs

这会触发严格化转换、函数内联、冗余计算消除等优化，显著降低惰性求值的开销。

4. 严格化自定义数据结构 #

使用BangPatterns扩展让Node字段严格求值，避免延迟计算：

{-# LANGUAGE BangPatterns #-}

data Heap a = Tip | Node !Rank !Int !(Heap Int) !(Heap Int)

同时修改makeHeap等函数，确保所有字段都被严格求值，减少thunk产生。

5. 维护总和避免遍历堆 #

在solve过程中维护当前总和，无需最后遍历整个堆计算：

solve :: Int -> Heap Int -> Int -> Int
solve 0 _ total = total
solve _ (Node _ 0 _ _) _ = 0
solve sips ppl total = solve (sips - 1) newHeap (total - val + newVal)
    where
        (val, h) = extract ppl
        newVal = val `div` 10
        newHeap = insert newVal h

main :: IO()
main = do
    line <- readFile "input.txt"
    let _:volume:ppl = map(read :: String -> Int) $ words line
    let heap = fromList ppl
    let initialTotal = sum ppl
    writeFile "output.txt" $ show $ solve volume heap initialTotal

此修改避免了toSum的递归遍历，节省O(n)的时间开销。

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