朴素Union-Find不相交集函数实现性能不佳原因咨询

咱先结合你给出的profiling数据和代码，拆解一下为啥这个朴素UFDS性能拉胯：

先看Profiling数据 #

total time = 0.10 secs (98 ticks @ 1000 us, 1 processor)
total alloc = 78,869,168 bytes (excludes profiling overheads)

COST CENTRE         MODULE  SRC                                      %time %alloc
x.\                 Main    src/merging_communities.hs:67:54-71       37.8   0.0
foldMap             Main    src/merging_communities.hs:(31,3)-(32,55) 22.4   0.0
x                   Main    src/merging_communities.hs:(65,1)-(68,79) 20.4  83.2
getElemTree         Main    src/merging_communities.hs:40:1-43        19.4   0.0
main.initialForest  Main    src/merging_communities.hs:103:7-51        0.0  16.2

你的代码实现 #

main.hs #

module Main where

import Control.Monad
import Control.Monad.State.Lazy
import Data.Foldable
import Data.Functor
import Data.List
import Data.Maybe
import Data.Monoid
import Prelude
import System.IO
import Text.Pretty.Simple

--The Union-Find algorithm and Disjoint Sets (UFDS) data structureare used which is able to efficiently (i.e. in nearly constant time) determine which set an item belongs to,
--test if two items belong to the same set, and union two disjoint sets into one when needed.
--It can be used to find connected components in an undirected graph, and can hence be used as part of Kruskal's algorithm for the Minimum Spanning Tree (MST) problem.

data Tree a = Node a [Tree a] deriving (Show)

instance (Eq a) => Eq (Tree a) where
  (Node a forestA) == (Node b forestB) = a == b && forestA == forestB

instance Functor Tree where
  fmap f (Node a []) = Node (f a) []
  fmap f (Node a forest) = Node (f a) (fmap (fmap f) forest)

instance Foldable Tree where
  foldMap f (Node a []) = f a
  foldMap f (Node a xs) = f a <> foldMap (foldMap f) xs

-- each disjoint set has a representative element which is used to uniquely identify the set. We can use a tree to represent a disjoint set where
-- the representative element is the root node of the tree
makeSet :: a -> Tree a
makeSet a = Node a []

getElemTree :: Eq a => a -> [Tree a] -> Maybe (Tree a)
getElemTree a forest = find (elem a) forest

size :: Tree a -> Int
size (Node a []) = 1
size (Node a forest) = 1 + (sum $ fmap size forest)

depth :: Tree a -> Int
depth (Node a []) = 1
depth (Node a forest) = 1 + (maximum $ fmap ((+ 1) . depth) forest)

flatten :: Tree a -> [a]
flatten (Node a forest) = [a] ++ (foldMap flatten forest)

-- set the parent of one of the roots to the other tree's root - which one we choose is based on our weighting
unWeightedUnion :: Eq a => a -> a -> [Tree a] -> [Tree a]
unWeightedUnion a b forest
  | isNothing treeA || treeA == treeB = forest
  | otherwise =
      let tA@(Node rootA forestA) = fromJust $ treeA
          tB@(Node rootB forestB) = fromJust $ treeB
      in changeRoot tA tB forest
  where
    treeA = getElemTree a forest
    treeB = getElemTree b forest
    changeRoot tA@(Node rootA forestA) tB@(Node rootB forestB) forest =
      if (size tA <= size tB)
        then (Node rootA (tB : forestA)) : filter (\t -> t /= tB && t /= tA) forest
        else (Node rootB (tA : forestB)) : filter (\t -> t /= tB && t /= tA) forest

-- union by rank is a weighting which keeps our trees as shallow as possible When we weight by rank or tree depth we make the shallower tree root the child of the deeper tree's root
getRoot :: Tree a -> a
-- get the root node
getRoot (Node a _) = a
-- return the name of the set containing the node x ie the root node of the set containing node x
-- use path compression - if parent is not the root then set the parent of the node to the root

data Query = M Int Int | Q Int deriving (Show, Read)

executeQuery :: [Query] -> Int -> StateT [Tree Int] IO Query
executeQuery [] _ = return $ M 1 1
executeQuery qs pop = do
  forest <- get
  case head qs of
    (M a b) -> do
      let newForest = unWeightedUnion a b forest
      put newForest
      executeQuery (tail qs) pop
    (Q a) -> do
      liftIO $ print $ size $ fromJust $ getElemTree a forest
      executeQuery (tail qs) pop

main = do
  contents <- readFile "queries.txt"
  print $ lines contents
  let population = head $ words contents
  let queries = map read $ tail $ lines contents :: [Query]
  let population = read $ head $ words contents :: Int
  let initialForest = map makeSet [1 .. population]
  execStateT (executeQuery queries population) initialForest

queries.txt #

100000 200000
M 68770 97917
M 65906 74478
M 78744 21384
M 36186 31560
Q 43063
M 12923 73331
M 91542 54702
M 62459 96133
M 13196 56121
M 1648 86052
M 99517 97247
M 59768 66017
Q 48274
Q 96430
M 44341 70873
Q 74989
Q 71357
M 72482 16677
Q 8219

核心性能瓶颈拆解 #

1. 线性查找森林中的树（getElemTree） #

getElemTree用find (elem a) forest来定位包含元素的树，这是**O(n)**的线性遍历。当你有10万初始元素时，每次Union/Query都要做两次这种查找，随着操作增多，虽然森林里的树数量减少，但每次查找的开销依然很大——profiling里它占了19.4%的时间就是明证。

2. Union操作的低效列表处理 #

unWeightedUnion里的filter (\t -> t /= tB && t /= tA) forest又是一次**O(n)**的线性遍历，而且每次都要构造新的树和列表。这直接导致了x这个cost centre占了83.2%的内存分配——每次Union都要复制整个森林列表的大部分元素，还会生成新的树节点，内存开销爆炸。

3. 完全没实现路径压缩 #

你代码里注释提到了路径压缩，但实际getRoot只是返回当前节点的根，根本没把路径上的节点直接挂到根上。没有路径压缩，树的深度会越来越大，后续的size计算、getElemTree查找都会因为树深增加而变慢，极端情况树会变成链状，每次操作都要遍历整个链。

4. size计算的递归遍历 #

每次调用size都要递归遍历整个树，这是**O(k)**的时间（k是树的节点数）。正常的UFDS应该维护每个集合的大小在一个单独的数组里，这样获取大小是O(1)的，而你每次Q查询都要遍历整个树，这在树很大时开销极大。

5. 用列表存储森林而非数组 #

列表的查找、修改都是O(n)的，而UFDS的标准实现是用数组（索引对应元素，值对应父节点），配合路径压缩和按秩合并，能做到近乎O(1)的操作时间。你用[Tree a]存森林，每次操作都要遍历整个列表，这在数据量大时完全不适用。

快速改进方向 #

• 改用数组存储父节点和大小：用两个数组，parent数组存每个元素的父节点，sizeArr数组存每个根节点对应的集合大小。
• 实现路径压缩：在find操作中，把路径上的所有节点直接挂到根节点上。
• 实现按秩合并：Union时把小的集合挂到大的集合下面，保持树的深度尽可能小。
• 去掉递归的树结构：完全抛弃Tree类型，用数组模拟父子关系，这才是UFDS高效的核心。

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