Haskell程序survivedeath时间复杂度不符预期的原因排查

问题代码 #

newtype Prob a = Prob { getProb :: [(a,Rational)] } deriving (Show,Eq,Functor)
flatten :: Prob (Prob a) -> Prob a
flatten (Prob xs) = Prob $ concat $ map multAll xs
    where multAll (Prob innerxs,p) = map (\(x,r) -> (x,p*r)) innerxs
instance Applicative Prob where 
    liftA2 fn (Prob x) (Prob y) = Prob [(fn a b,prob1 *prob2) |(a,prob1) <- x, (b,prob2) <- y]
    pure x = Prob [(x,1%1)]
instance Monad Prob where
    m >>= f = flatten (fmap f m)
makedie x = Prob (zip [1..x] (repeat (1%x)))
proboperate2 op a b= sumprobs (liftA2 op a b)
rerolldie op = proboperate2 op <*> id
sumprobs (Prob a) = Prob [(v,sum (map snd (filter ((==v) . fst) a))) |v <- indivalues]
    where indivalues = nub (map fst a)
survivedeath :: Integer -> Prob Integer -> Prob Bool 
survivedeath dc die = sumprobs (survivegiven (0,0) =<< die) 
    where 
        survivegiven :: (Integer,Integer) -> Integer -> Prob Bool
        survivegiven (a,_) _ | a >= 3 = return False 
        survivegiven (_,a) _ | a >= 3 = return True 
        survivegiven (a,b) 1 = sumprobs ((survivegiven (a+2,b)) =<< die)
        survivegiven (a,b) 20 = return True 
        survivegiven (a,b) n | n >= dc = sumprobs ((survivegiven (a,1+b)) =<< die)
        survivegiven (a,b) n  = sumprobs ((survivegiven (1+a,b)) =<< die)

性能疑问 #

这段代码中的survivedeath函数运行速度会迅速变慢。我原本认为sumprobs的时间复杂度是O(N²)，>>=是O(N)，针对d20运行6次应该只有约2400次操作，但实际运行并不快，请问这是为什么？

原因分析 #

1. 指数级的列表膨胀
   你的操作次数估算完全错误：每次>>=（Monad绑定）操作在未合并重复项时，会让列表长度呈指数级增长。比如用d20时：

   • 第一次绑定后列表长度是20
   • 第二次是20×20=400
   • 第三次是20³=8000
   • 第六次会达到20⁶=64,000,000项
     此时sumprobs处理这个规模的列表，O(N²)的复杂度意味着要执行约4×10¹⁵次操作，这显然不可能快速完成。

2. sumprobs的低效实现
   sumprobs里的nub和filter都是O(N)操作，且每次都要遍历整个列表去重、求和。当列表长度达到几十万甚至百万级时，这些操作的开销会被放大到无法接受的程度。

3. 递归中的重复sumprobs调用
   在survivegiven的递归分支里，每次绑定后都立刻调用sumprobs，虽然能减少后续绑定的列表长度，但每次sumprobs本身的开销已经随着递归深度增加而急剧上升。

优化建议 #

最核心的优化是用高效的数据结构替代列表存储概率分布：

• 改用Data.Map（或Data.IntMap如果键是整数）来存储(值,概率)对，Map会自动对键去重，且插入、查找、合并操作的时间复杂度都是O(logN)。
• 重新定义Prob类型为：

  import qualified Data.Map as Map
  newtype Prob a = Prob { getProb :: Map.Map a Rational } deriving (Show, Eq, Functor)
  

• 此时sumprobs可以简化为恒等函数（因为Map已经自动合并相同键的概率），>>=和liftA2操作可以基于Map的折叠、合并实现，避免指数级的列表膨胀。

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