使用python實現二叉搜索樹

準備

我們想要創建一個二叉搜索樹結構來讓所有的元數據都按正確的數據存儲，同時我們希望這是一個可變的結構，所以主要有以下幾點：

1. 設計基本的二叉搜索樹結構

2. 使用MutableSet結構作爲基類。

   具體的介紹可參照python官方collections模塊docs

3. 二叉搜索樹主要有兩個分支：一個分支用於存放小於當前節點的鍵，另一個用於存放大於當前節點的鍵。具體的二叉搜索樹可以參照百度文庫。所以我們需要研究如何將集合與抽象基類集成。

4. 這不會是一個龐大的序列。通常使用鍵來引用對應的元素。

5. 可以讓所有的鍵按順序存儲，這也正是二叉搜索樹樹的一個用途。

設計

將分爲兩個類： Tree 和 TreeNode。

TreeNode作爲樹的節點，包含more(大於的),less(小於的)，parent(父節點)的引用。同時爲了搜索一個指定的元素，會使用一個遞歸操作將搜索任務委託給節點自身來完成：

1. 如果目標元素與當前元素相當，那麼返回self
2. 如果目標元素比當前元素小，那麼遞歸使用less.find。即在“小於分支”遞歸搜索
3. 如果目標元素比當前元素打，那麼遞歸使用more.find

Tree主要爲MutableSet抽象基類提供必須的接口。

Tree類

from collections import MutableSet
import weakref

class Tree(MutableSet):
    # 根節點沒有值，即None
    def __init__(self, iterable=None):
        self.root = TreeNode(None)  # 代表根節點
        self.size = 0				# 這個樹結構的大小
        if iterable:				# 接受一個可迭代對象，並加載對象的所有元素
            for item in iterable:
                self.root.add(item)

    def add(self, item):
        # 每有元素添加，當將大小加1，這樣當我們需要知道當前節點總數時就不需要遍歷了。
        self.root.add(item)
        self.size += 1

    def discard(self, item):
        # 每有元素添加，當將大小減1，這樣當我們需要知道當前節點總數時就不需要遍歷了。
        try:
            self.root.remove(item)
            self.size -= 1
        except KeyError:
            pass

    def __contains__(self, item):
        # 如果當前結構包含item即返回Treu 否則返回False
        try:
            self.root.more.find(item)
            return True
        except KeyError:
            return False

    def __iter__(self):
        # 這個函數是控制生成器的函數。
        # root節點爲None，所以後續的節點在more分支上。
        for item in iter(self.root.more):
            yield item

    def __len__(self):
        # 這個函數是控制len(obj)的
        return self.size

TreeNode類

class TreeNode:
    def __init__(self, item, less=None, more=None, parent=None):
        self.item = item 			# 代表當前節點 
        self.less = less			# 小於當前節點的分支
        self.more = more			# 大於當前節點的分支
        if parent is not None:
            self.parent = parent

    @property
    def parent(self):
        return self.parent_ref()

    @parent.setter
    def parent(self,value):
        # 使用弱引用，方便內存回收。
        self.parent_ref = weakref.ref(value)

    def __repr__(self):
        return ("TreeNode({item!r},{less!r},{more!r})".format(**self.__dict__))

    def find(self, item):
        if self.item is None:
            # 爲None時代表當前爲根節點，所以向more分支遞歸搜索。
            if self.more:
                return self.more.find(item)
        elif self.item == item:
            # 相等時爲找到了節點
            return self
        elif self.item > item and self.less:
            # 如果當前節點大於目標對象，並且小於分支存在的話，那麼向less分支遞歸搜索
            return self.less.find(item)
        elif self.item < item and self.more:
            # 如果當前節點小於目標對象，並且大於分支存在的話，那麼向more分支遞歸搜索
            return self.more.find(item)
        # 如果以上判斷都不符合條件，那麼搜索的元素不存在
        raise KeyError

    def __iter__(self):
        # 如果當前節點小於分支存在，那麼遍歷並返回。
        # 使用yield可以節省開銷。
        if self.less:
            for item in iter(self.less):
                yield item
        # 返回當前節點
        yield self.item
        # 如果當前節點大於分支存在，那麼遍歷並返回。
        if self.more:
            for item in iter(self.more):
                yield item

    def add(self, item):
        # 爲None時代表當前爲根節點，向more分支添加節點
        if self.item is None:
            if self.more:
                self.more.add(item)
            else:
                self.more = TreeNode(item, parent=self)
        # 如果當前節點大於等於添加節點，那麼向less分支添加節點
        elif self.item >= item:
            if self.less:
                self.less.add(item)
            else:
                self.less = TreeNode(item, parent=self)
        # 如果當前節點小於添加節點，那麼向more節點
        elif self.item < item:
            if self.more:
                self.more.add(item)
            else:
                self.more = TreeNode(item, parent=self)

    def remove(self, item):
        # 如果當前節點爲None或者目標節點大於當前節點，那麼向more分支遞歸
        if self.item is None or item> self.item:
            if self.more:
                self.more.remove(item)
            else:
                raise KeyError
        # 如果目標節點小於當前節點，那麼向less分支遞歸
        elif item < self.item:
            if self.less:
                self.less.remove(item)
            else:
                raise KeyError
        else: # self.item == item 即找到了目標元素
            # 如果當前節點具有less分支和more分支
            if self.less and self.more:
                successor = self.more._least() 	# 遞歸找當前節點more分支中的最小節點
                self.item = successor.item  	# 並將當前節點的值設置爲最小節點的值
                successor.remove(successor.item) # 繼續遞歸尋找合適的_replace操作
            # 如果當前節點僅有less分支
            elif self.less:
                self._replace(self.less)
            elif self.more:
                self._replace(self.more)
            # 葉子節點
            else:
                self._replace(None)

    def _least(self):
        # 遞歸搜索最小的節點
        if self.less is None:
            return self
        return self.less._least()

    def _replace(self,new=None):
        # 如果當前節點存在父節點
        if self.parent:
            # 如果當前節點在父節點的小於分支，那麼將小於分支設置爲new值
            if self == self.parent.less:
                self.parent.less = new
            # 如果當前節點在父節點的大於分支，那麼將大於分支設置爲new值
            else:
                self.parent.more = new
        # 如果指定了父節點，那麼替換父節點指向
        if new is not None:
            new.parent = self.parent

remove函數的過程會晦澀一些，大致如下：

1. 如果當前節點爲None或者目標節點大於當前節點，那麼向more分支遞歸

2. 如果目標節點小於當前節點，那麼向less分支遞歸

3. 如果 self.item == item 即當前節點爲目標節點，進行後續操作

   1. 如果當前節點同時存在less分支和more分支，那麼尋找more分支中的最小節點，將當前節點的值設爲最小節點的值，並繼續遞歸尋找合適的_replace操作。

      這裏之所以尋找more分支中的最小節點，是因爲二叉搜索樹的結構。more分支的所有節點必定都比less分支大，所以替換完成後不影響順序

   2. 如果當前節點只有less分支，或這more分支，那麼直接執行_replace中的邏輯。

使用

>> a = Tree([1,6,3,4,8,7,9,2])
>> list(a)
Out[4]: [1, 2, 3, 4, 6, 7, 8, 9]
>> len(a)
Out[5]: 8
>> a.remove(8)
>> list(a)
Out[9]: [1, 2, 3, 4, 6, 7, 9]

可以看到，數據始終時有序的。