''' mbinary ######################################################################### # File : bTree.py # Author: mbinary # Mail: zhuheqin1@gmail.com # Blog: https://mbinary.xyz # Github: https://github.com/mbinary # Created Time: 2018-08-29 12:49 # Description: ######################################################################### ''' class node: def __init__(self, keys=None, isLeaf=True, children=None): if keys is None: keys = [] if children is None: children = [] self.keys = keys self.isLeaf = isLeaf self.children = [] def __getitem__(self, i): return self.keys[i] def __delitem__(self, i): del self.keys[i] def __setitem__(self, i, k): self.keys[i] = k def __len__(self): return len(self.keys) def __repr__(self): return str(self.keys) def __str__(self): children = ','.join([str(nd.keys) for nd in self.children]) return f'keys: {self.keys}\nchildren: {children}\nisLeaf: {self.isLeaf}' def getChd(self, i): return self.children[i] def delChd(self, i): del self.children[i] def setChd(self, i, chd): self.children[i] = chd def getChildren(self, begin=0, end=None): if end is None: return self.children[begin:] return self.children[begin:end] def findKey(self, key): for i, k in enumerate(self.keys): if k >= key: return i return len(self) def update(self, keys=None, isLeaf=None, children=None): if keys is not None: self.keys = keys if children is not None: self.children = children if isLeaf is not None: self.isLeaf = isLeaf def insert(self, i, key=None, nd=None): if key is not None: self.keys.insert(i, key) if not self.isLeaf and nd is not None: self.children.insert(i, nd) def isLeafNode(self): return self.isLeaf def split(self, prt, t): # form new two nodes k = self[t-1] nd1 = node() nd2 = node() # note that t is 1 bigger than key index nd1.keys, nd2.keys = self[:t-1], self[t:] nd1.isLeaf = nd2.isLeaf = self.isLeaf if not self.isLeaf: # note that children index is one bigger than key index, and all children included nd1.children, nd2.children = self.children[0:t], self.children[t:] # connect them to parent idx = prt.findKey(k) if prt.children != []: prt.children.remove(self) # remove the original node prt.insert(idx, k, nd2) prt.insert(idx, nd=nd1) return prt class bTree: def __init__(self, degree=2): self.root = node() self.degree = degree self.nodeNum = 1 self.keyNum = 0 def search(self, key, withpath=False): nd = self.root fathers = [] while True: i = nd.findKey(key) if i == len(nd): fathers.append((nd, i-1, i)) else: fathers.append((nd, i, i)) if i < len(nd) and nd[i] == key: if withpath: return nd, i, fathers else: return nd, i if nd.isLeafNode(): if withpath: return None, None, None else: return None, None nd = nd.getChd(i) def insert(self, key): if len(self.root) == self.degree*2-1: self.root = self.root.split(node(isLeaf=False), self.degree) self.nodeNum += 2 nd = self.root while True: idx = nd.findKey(key) if idx < len(nd) and nd[idx] == key: return if nd.isLeafNode(): nd.insert(idx, key) self.keyNum += 1 return else: chd = nd.getChd(idx) # ensure its keys won't excess when its chd split and u if len(chd) == self.degree*2-1: nd = chd.split(nd, self.degree) self.nodeNum += 1 else: nd = chd def delete(self, key): # to do '''search the key, delete it , and form down to up to rebalance it ''' nd, idx, fathers = self.search(key, withpath=True) if nd is None: return del nd[idx] self.keyNum -= 1 if not nd.isLeafNode(): chd = nd.getChd(idx) # find the predecessor key while not chd.isLeafNode(): fathers.append((chd, len(chd)-1, len(chd))) chd = chd.getChd(-1) fathers.append((chd, len(chd)-1, len(chd))) nd.insert(idx, chd[-1]) del chd[-1] if len(fathers) > 1: self.rebalance(fathers) def rebalance(self, fathers): nd, keyIdx, chdIdx = fathers.pop() while len(nd) < self.degree-1: # rebalance tree from down to up prt, keyIdx, chdIdx = fathers[-1] lbro = [] if chdIdx == 0 else prt.getChd(chdIdx-1) rbro = [] if chdIdx == len(prt) else prt.getChd(chdIdx+1) if len(lbro) < self.degree and len(rbro) < self.degree: # merge two deficient nodes beforeNode, afterNode = None, None if lbro == []: keyIdx = chdIdx beforeNode, afterNode = nd, rbro else: beforeNode, afterNode = lbro, nd keyIdx = chdIdx-1 # important, when choosing keys = beforeNode[:]+[prt[keyIdx]]+afterNode[:] children = beforeNode.getChildren() + afterNode.getChildren() isLeaf = beforeNode.isLeafNode() prt.delChd(keyIdx+1) del prt[keyIdx] nd.update(keys, isLeaf, children) prt.children[keyIdx] = nd self.nodeNum -= 1 elif len(lbro) >= self.degree: # rotate when only one sibling is deficient keyIdx = chdIdx-1 nd.insert(0, prt[keyIdx]) # rotate keys prt[keyIdx] = lbro[-1] del lbro[-1] if not nd.isLeafNode(): # if not leaf, move children nd.insert(0, nd=lbro.getChd(-1)) lbro.delChd(-1) else: keyIdx = chdIdx nd.insert(len(nd), prt[keyIdx]) # rotate keys prt[keyIdx] = rbro[0] del rbro[0] if not nd.isLeafNode(): # if not leaf, move children # note that insert(-1,ele) will make the ele be the last second one nd.insert(len(nd), nd=rbro.getChd(0)) rbro.delChd(0) if len(fathers) == 1: if len(self.root) == 0: self.root = nd self.nodeNum -= 1 break nd, i, j = fathers.pop() def __str__(self): head = '\n'+'-'*30+'B Tree'+'-'*30 tail = '-'*30+'the end'+'-'*30+'\n' lst = [[head], [f'node num: {self.nodeNum}, key num: {self.keyNum}']] cur = [] ndNum = 0 ndTotal = 1 que = [self.root] while que != []: nd = que.pop(0) cur.append(repr(nd)) ndNum += 1 que += nd.getChildren() if ndNum == ndTotal: lst.append(cur) cur = [] ndNum = 0 ndTotal = len(que) lst.append([tail]) lst = [','.join(li) for li in lst] return '\n'.join(lst) def __iter__(self, nd=None): if nd is None: nd = self.root que = [nd] while que != []: nd = que.pop(0) yield nd if nd.isLeafNode(): continue for i in range(len(nd)+1): que.append(nd.getChd(i)) if __name__ == '__main__': bt = bTree() from random import shuffle, sample n = 20 lst = [i for i in range(n)] shuffle(lst) test = sample(lst, len(lst)//4) print(f'building b-tree with {lst}') for i in lst: bt.insert(i) # print(f'inserting {i}) # print(bt) print(bt) print(f'serching {test}') for i in test: nd, idx = bt.search(i) print(f'node: {repr(nd)}[{idx}]== {i}') for i in test: print(f'deleting {i}') bt.delete(i) print(bt)