在代碼設(shè)計中時常面對這樣的場景，給定兩個元素，我們需要快速判斷他們是否屬于同一個集合，同時不同的集合在需要時還能快速合并為一個集合，例如我們要開發(fā)一個社交應(yīng)用，那么判斷兩個用戶是否是朋友關(guān)系，或者兩人是否屬于同一個群就需要用到我們現(xiàn)在提到的功能。

這些功能看似簡單，但有個難點在于你要處理的“足夠快”，假設(shè)a,b兩個元素分別屬于集合A,B,判斷它們是否屬于同一個集合的直接做法就是遍歷集合A中所有元素，看看是否能找到b，如果集合A中包含n個元素，那么該做法的時間復(fù)雜度就是O(n)，當集合元素很多，而且判斷的次數(shù)也很多時，這樣的做法效率就會很低，本節(jié)我們要看看能不能找到次線性的算法。

我們先看復(fù)雜度為O(n)的算法邏輯，假設(shè)我們有6個元素，編號分別為0到6，我們可以使用隊列來模擬集合，屬于同一個集合的元素就存儲在同一個隊列中，然后每個元素通過哈希表映射到隊列頭，如下圖所示：

在這種數(shù)據(jù)結(jié)構(gòu)下，查詢兩個元素是否屬于同一個集合，那么只要通過哈希表找到各自元素所在隊列的頭部，判斷頭部是否一致即可，我們用areDisjoint(x,y)來表示兩個元素是否屬于一個集合，那么在當前數(shù)據(jù)結(jié)構(gòu)下areDisjoint的時間復(fù)雜度是O(1)。

如果要合并兩個元素所在集合，我們用merge(x,y)來表示，那么在當前結(jié)構(gòu)下，我們只要找到x和y對應(yīng)的隊列頭部，然后從x所在隊列的頭部遍歷到最后一個元素，然后將最后一個元素的next指針執(zhí)行y所在的隊列頭部，如下圖所示：

同時我們還需要做一個操作，那就是修改第二個集合中每個元素映射的隊列頭部，因此在當前結(jié)構(gòu)下，merge(x,y)對應(yīng)時間復(fù)雜度為O(n)，因為從隊列頭遍歷到末尾是O(n),同時遍歷y所在集合每個元素，修改他們映射的隊列頭，時間復(fù)雜度也是O(n)。

現(xiàn)在問題是我們能否將合并所需要的時間進行優(yōu)化。我們注意到合并時有兩個步驟很耗時，一是從隊列走到隊尾，二是修改第二個集合中每個元素指向的隊列頭。所以耗時其實是因為我們使用隊列來表示集合所導(dǎo)致。為了優(yōu)化時間，我們將隊列換成多叉樹，如下圖所示：

此時我們不再使用哈希表來將元素映射到隊列頭部，而是將同一個集合的元素安插到同一個多叉樹中，要判斷兩個元素是否屬于同一集合，我們只要沿著元素的父節(jié)點指針往上走一直找到樹的根節(jié)點，如果找到相同的根節(jié)點，那么兩個元素就屬于同一集合，對于排序二叉樹而言，樹的高度為O(lg(n))，n是樹的節(jié)點數(shù)，于是判斷兩個元素是否屬于同一集合所需時間復(fù)雜度為O(lg(n))。

當需要合并兩個元素對于的集合時，我們分別找到兩個元素對于的根節(jié)點，然后將高度較低的那棵樹的根節(jié)點作為高度較高那棵樹的子節(jié)點，這個處理對效率很重要，后面我們會進一步研究，樹合并的情形如下圖所示：

下面我們先看看代碼實現(xiàn)：

# This is a sample Python script. 
 
# Press ⌃R to execute it or replace it with your code. 
# Press Double ⇧ to search everywhere for classes, files, tool windows, actions, and settings. 
 
class Element: 
    def __init__(self, val : int): 
        self.val = val 
        self.parent = self #元素在創(chuàng)建時自己形成一個單獨集合，因此父節(jié)點指向自己 
    def value(self): 
        return self.val 
    def parent(self): 
        return self.parent 
    def set_parent(self, parent): 
        assert parent is not None 
        self.parent = parent 
 
class DisjontSet: 
    def __init__(self): 
        self.hash_map = {} 
    def add(self, elem : Element): 
        assert elem is not None 
        if elem.value() in self.hash_map: 
            return False 
        self.hash_map[elem.value()] = elem 
        return True 
    def find_partition(self, elem : Element): 
        #返回元素所在集合的根節(jié)點 
        assert elem is not None or elem.value() in self.hash_map 
        parent = elem.parent() 
        if parent != elem: #遞歸查找根節(jié)點,樹的高度為lg(n)，所以這里查找的時間復(fù)雜度為lg(n) 
            parent = self.find_partition(parent) 
        return parent 
 
    def are_disjoint(self, elem1 : Element, elem2 : Element): 
        #判斷兩個元素是否屬于同一集合只要判斷他們再哈希表中映射的根節(jié)點是否同一個 
        root1 = self.find_partition(elem1) 
        root2 = self.find_partition(elem2) 
        return root1 is not root2 
 
    def merge(self, elem1 : Element, elem2 : Element): 
        root1 = self.find_partition(elem1) 
        root2 = self.find_partition(elem2) 
        if root1 is root2: 
            #兩個元素屬于同一個集合 
            return False 
        root2.setParent(root1) 
        self.hash_map[root2.value()] = root1 #設(shè)置root2對應(yīng)的父節(jié)點 
 
# Press the green button in the gutter to run the script. 
if __name__ == '__main__': 
 
# See PyCharm help at https://www.jetbrains.com/help/pycharm/ 

由于我們將集合的表示從隊列改為了多叉樹，因此集合的查找與合并對應(yīng)復(fù)雜度為O(lg(n))，現(xiàn)在問題是我們能否繼續(xù)改進效率。當前merge函數(shù)耗時在于我們要通過parent指針一直爬到根節(jié)點，如果能讓parent指針直接指向根節(jié)點那么不就省卻向上爬的時間開銷嗎，這種直接將下層節(jié)點父指針直接指向根節(jié)點的辦法叫路徑壓縮，如下圖所示：

從上圖看到，節(jié)點6，8的父節(jié)點原來是9，它所在集合的根節(jié)點是1，于是我們直接將原來指向9的指針直接指向根節(jié)點1，這樣以后在合并或查詢集合時我們就可以省掉向上爬的時間開銷。還有一個問題在上面代碼中兩棵樹合并問題，我們僅僅是把root2的父指針指向root1，這么做會存在合并后樹不平衡問題，也就是合并后的左右子樹高度可能相差較大，這種情況也會對效率產(chǎn)生不利影響，如下圖所示：

可以看到右下角合并后左右子樹高度差異大，于是節(jié)點，6，8找到根節(jié)點0所需的時間就要比2，3，4要多，但形成右上角的情況時，葉子節(jié)點6，8和2，3，4找到根節(jié)點的時間就差不多，這樣就有利于效率的提高，所以我們還需要記錄下樹的高度，在合并時要將高度小的樹合向高度高的樹，因此代碼修改如下：

class Element: 
    def __init__(self, val : int): 
        self.val = val 
        self.parent = self #元素在創(chuàng)建時自己形成一個單獨集合，因此父節(jié)點指向自己 
        self.rank = 1 #表示樹的高度 
    def value(self): 
        return self.val 
    def parent(self): 
        return self.parent 
    def set_parent(self, parent): 
        assert parent is not None 
        self.parent = parent 
    def get_rank(self): 
        return self.rank  
    def set_rank(self, rank): 
        assert rank > 1 
        self.rank = rank 

然后我們需要修改find_partition的做法

def find_partition(self, elem : Element): 
    #返回元素所在集合的根節(jié)點 
    assert elem is not None or elem.value() in self.hash_map 
    parent = elem.parent() 
    if parent is elem: #已經(jīng)是根節(jié)點 
        return elem  
    parent = self.find_partition(elem) #獲得集合的根節(jié)點 
    elem.set_parent(parent) #路徑壓縮直接指向根節(jié)點 
    return parent #返回根節(jié)點 

注意到find_partion的實現(xiàn)中有遞歸過程，如果當前節(jié)點不是根節(jié)點，那么遞歸的查詢根節(jié)點，然后把當前節(jié)點的parent指針直接指向根節(jié)點，我們看到這步修改所需的時間復(fù)雜度跟原來一樣都是lg(n)。

接下來我們要修改merge的實現(xiàn)：

def merge(self, elem1 : Element, elem2 : Element): 
        root1 = self.find_partition(elem1) 
        root2 = self.find_partition(elem2) 
        if root1 is root2:  # 兩個元素屬于同一個集合 
            return False 
        new_rank = root1.get_rank() + root2.get_rank() 
        if root1.get_rank() >= root2.get_rank():  # 根據(jù)樹的高度來決定合并方向 
            root2.set_parent(root1) 
            root1.set_rank(new_rank) 
        else: 
            root1.set_parent(root2) 
            root2.set_rank(new_rank) 
        return True 

這種改進后，在m次指向find_partion和merge調(diào)用時所需要的時間是O(m)，也就是說在改進后，當大量調(diào)用find_partion和merge時，這些調(diào)用的平均耗時降到了O(1)，也就是說路徑壓縮后，其效果在大批量的調(diào)用查找集合和合并集合操作時能出現(xiàn)非常顯著的效率提升，其對應(yīng)的數(shù)學(xué)證明非常負責，我們暫時忽略調(diào)。我們可能對這里的效率提升感受不到，但想想微信中對兩個人是否屬于同一個群的調(diào)用一天至少也有千萬乃至上億次吧，因此這里的改進能大大的改進服務(wù)器的處理效率。

完整代碼在這里

https://github.com/wycl16514/python_disjoint_set.git