yawei 2020-06-12
class BSTMapNode(object): def __init__(self, key, value): self.key = key self.value = value self.left = None self.right = None # 以列表作为底层存储 class BSTMapIterator(object): def __init__(self, root): self.the_keys = list() self.cur_item = 0 # 初始化内置list,并将索引reset self.bst_travel(root) self.cur_item = 0 def __iter__(self): return self # 遍历列表即可 def __next__(self): if self.cur_item < len(self.the_keys): key = self.the_keys[self.cur_item] self.cur_item += 1 return key else: raise StopIteration # 这个递归!中序遍历 def bst_travel(self, sub_tree): if sub_tree is None: return self.bst_travel(sub_tree.left) self.the_keys[self.cur_item] = sub_tree.key self.cur_item += 1 self.bst_travel(sub_tree.right) class BSTMap(object): def __init__(self): self._root = None self._size = None def __len__(self): return self._size def __iter__(self): return BSTMapIterator(self._root) def __contains__(self, key): return self.bst_search(self._root, key) is not None def value_of(self, key): node = self.bst_search(self._root, key) return node.value # 递归搜索 def bst_search(self, sub_tree, target): if sub_tree is None: return None elif target < sub_tree.key: return self.bst_search(sub_tree.left, target) elif target > sub_tree.key: return self.bst_search(sub_tree.right, target) else: return sub_tree # 查找最小节点 def bst_min(self, sub_tree): if sub_tree is None: return None elif sub_tree.left is None: return sub_tree else: return self.bst_min(sub_tree.left) # 若key=sub_tree的key ,替换value def add(self, sub_tree, key, value): node = self.bst_search(sub_tree, key) if node is not None: node.value = value return False else: self._root = self._bst_insert(self._root, key, value) self._size += 1 return True # 如果树为空,则将新建节点并返回,若key小于根节点,则递归插入左子树,key大于根节点,递归插入右子树 def _bst_insert(self, sub_tree, key, value): if sub_tree is None: sub_tree = BSTMapNode(key, value) elif key < sub_tree.key: sub_tree.left = self._bst_insert(sub_tree.left, key, value) elif key > sub_tree.key: sub_tree.right = self._bst_insert(sub_tree.right, key, value) return sub_tree def remove(self, key): self._root = self._bst_remove(self._root, key) self._size -= 1 def _bst_remove(self, sub_tree, target): if sub_tree is None: return sub_tree # 目标比当前节点小,则递归删除左子树的对应节点,否则递归删除右子树对应节点 elif target < sub_tree.key: sub_tree.left = self._bst_remove(sub_tree.left, target) return sub_tree elif target > sub_tree.key: sub_tree.right = self._bst_remove(sub_tree.right, target) return sub_tree # 目标等于当前节点 else: # 目标是叶子节点 if sub_tree.left is None and sub_tree.right is None: sub_tree = None # 目标只有左子树或右子树 elif sub_tree.left is None or sub_tree.right is None: if sub_tree.left is not None: sub_tree = sub_tree.left else: sub_tree = sub_tree.right # 目标同时有左右子树 # 查找右子树的最小节点,置为当前节点,并将其删除 else: successor = self.bst_min(sub_tree.right) sub_tree.key = successor.key sub_tree.value = successor.value sub_tree.right = self._bst_remove(sub_tree.right, successor.key) return sub_tree