""" Python 2-3 Tree implementation 2-3 Tree is a balanced tree each node of which may contain 2 elements and 3 references on its children. Element lookup speed is log2(N) < x < log3(N) Insertion and deletion is about 2 * log2(N) See http://en.wikipedia.org/wiki/2-3_tree for more info 2011 by Boris Tatarintsev """ class Pair(object): # use this class if associative tree (or map) is needed # over 2-3 tree def __init__(self, key, value): self.key = key self.value = value def __lt__(self, other): if type(other) is Pair: return self.key < other.key else: return self.key < other def __gt__(self, other): if type(other) is Pair: return self.key > other.key else: return self.key > other def __eq__(self, other): if type(other) is Pair: return self.key == other.key else: return self.key == other return None def __str__(self): return 'key: %s, value: %s' % (str(self.key), str(self.value)) def key(self): return self.key def val(self): return self.value class Node(object): def __init__(self, v = None, parent = None): self.values, self.valcnt = None, 0 self.links, self.refcnt = None, 0 self.parent = parent self.insertValue(v) def __str__(self): out = [] if self.values is not None: for v in self.values: if v is not None: out.append(' %s ' % str(v)) return ''.join(out) else: return 'empty' def __iter__(self): if self.values is not None: for item in self.values: yield item def __getlink(self, a): for idx in xrange(self.valcnt): if idx is 0: if a < self.values[idx]: return idx else: if self.values[idx - 1] < a < self.values[idx]: return idx if idx == self.valcnt - 1: return idx + 1 return -1 def __addLink(self, link): if self.links is None: self.links = [None] * 4 self.links[self.refcnt] = link self.refcnt += 1 def __insertLink(self, idx, anotherNode): if self.links is None: self.links = [None] * 4 if idx == 0: self.links[0],self.links[1],self.links[2], self.links[3] = anotherNode,self.links[0],self.links[1], self.links[2] elif idx == 1: self.links[1], self.links[2], self.links[3] = anotherNode, self.links[1], self.links[2] elif idx == 2: self.links[2], self.links[3] = anotherNode, self.links[2] else: self.links[3] = anotherNode self.refcnt += 1 def __removeLink(self, idx): if idx == 0: self.links[0], self.links[1], self.links[2], self.links[3] = self.links[1], self.links[2], self.links[3], None elif idx == 1: self.links[1], self.links[2], self.links[3] = self.links[2], self.links[3], None elif idx == 2: self.links[2], self.links[3] = self.links[3], None else: self.links[3] = None self.refcnt -= 1 def __rearrangeLinks(self, a): """ Rearrange links when adding a new node """ if self.valcnt != 0: if a < self.values[0] and not self.isLeafNode() and self.refcnt < 3: # shift all the links to the right when adding new in element self.__insertLink(0, None) elif self.valcnt == 2 and self.refcnt == 3 and self.values[self.valcnt - 1] > a > self.values[0]: # rearrange middle links when adding med element self.__insertLink(1, None) def __sort3(self, arr, l): """ Sort 2 or 3 arrays (very rubost and fast) """ if l >= 2: if arr[0] > arr[1]: arr[0], arr[1] = arr[1], arr[0] if l == 3: if arr[1] > arr[2]: arr[1], arr[2] = arr[2], arr[1] if arr[0] > arr[1]: arr[0], arr[1] = arr[1], arr[0] # interface methods & properties def insertValue(self, a): """ Insert a value into node """ if a is not None and self.valcnt < 3: if self.valcnt is 0: self.values = [None] * 3 self.__rearrangeLinks(a) self.values[self.valcnt] = a self.valcnt += 1 self.__sort3(self.values, self.valcnt) return self def removeValue(self, val): """ Remove value from node """ if self.contains(val): idx = self.values.index(val) if idx == 0: self.values[0], self.values[1], self.values[2] = self.values[1], self.values[2], None elif idx == 1: self.values[1], self.values[2] = self.values[2], None else: self.values[2] = None self.valcnt -= 1 return self def removeLink(self, node): """ Remove link from self to another node """ self.__removeLink(self.getLinkIdx(node)) return self def isConsistent(self): """ Check whether the node is consistent, this means it doesn't contain 3 items or 4 links """ return not (self.valcnt > 2 or self.refcnt > 3) def isLeafNode(self): """ Check whether this is a leaf node or not """ return self.refcnt == 0 def isEmptyNode(self): """ Returns true if node doesn't containt any value """ return self.valcnt == 0 def getLink(self, linkIdx): """ Get link by its index, return None if there is no link with such an index """ if linkIdx < self.refcnt: return self.links[linkIdx] def getLinkIdx(self, destNode): """ Get index of the link which points to the given node """ return self.links.index(destNode) def addLink(self, anotherNode): """ Add link to another node """ if anotherNode is not None: if self.links is None: self.links = [None] * 4 idx = self.__getlink(anotherNode.values[0]) if idx != -1: if idx < self.refcnt and self.links[idx] is None: self.links[idx] = anotherNode else: self.__insertLink(idx, anotherNode) anotherNode.parent = self return self def contains(self, a): """ Check if node contains a given value """ if self.valcnt is not 0: if (self.values[0] > a or self.values[self.valcnt - 1] < a) or a not in self.values: return None return self.values[self.values.index(a)] def chooseChild(self, a): """ Choose where to go according to the value a """ idx = self.__getlink(a) if 0 <= idx < self.refcnt: return self.links[idx] def getItem(self, a): if self.contains(a): return self.values[self.values.index(a)] return None class TTTree(object): def __init__(self): self.root = Node() self.lastSearchDepth = 0 def __iter__(self): stack = [self.root] while len(stack): node = stack.pop() yield node for j in xrange(node.refcnt - 1, -1, -1): stack.append(node.getLink(j)) def __str__(self): """ String representation of a tree (parentheses form) """ out, stack = [], [self.root] while stack: node = stack.pop() if node == ')': out.append(')') continue out.append('%s(' % str(node)) stack.append(')') for j in xrange(node.refcnt - 1, -1, -1): stack.append(node.getLink(j)) return ''.join(out) def __nextSucc(self, node): self.lastSearchDepth += 1 if not node.isLeafNode(): return self.__nextSucc(node.links[0]) return node def __find(self, curNode, a): if curNode is not None: if curNode.contains(a): return curNode nextNode = curNode.chooseChild(a) if nextNode is None: return curNode self.lastSearchDepth += 1 return self.__find(nextNode, a) def __getLeftSibling(self, node): """ Returns left sibling of a node """ if (node and node.parent) is not None: return node.parent.getLink(node.parent.getLinkIdx(node) - 1) def __getRightSibling(self, node): """ Returns right sibling of a node """ if (node and node.parent) is not None: return node.parent.getLink(node.parent.getLinkIdx(node) + 1) def __getSiblings(self, node): """ Returns node's siblings """ # check whether one of our siblings has two items lS, rS, lCnt, rCnt = None, None, 0, 0 if self.__getRightSibling(node) is not None: rS = self.__getRightSibling(node) rCnt = rS.valcnt if self.__getLeftSibling(node) is not None: lS = self.__getLeftSibling(node) lCnt = lS.valcnt return lS, lCnt, rS, rCnt def __swapValues(self, node1, a1, node2, a2): """ Swap any two values in nodes """ if node1 is not node2: idx1, idx2 = node1.values.index(a1), node2.values.index(a2) node1.values[idx1], node2.values[idx2] = node2.values[idx2], node1.values[idx1] def __fixNodeRemove(self, node, parent = -1): """ Fix deletion """ if node.isEmptyNode(): if node is not self.root: if parent == -1: parent = node.parent if node.isEmptyNode() or not node.isConsistent(): lS, lCnt, rS, rCnt = self.__getSiblings(node) rSS, lSS = self.__getRightSibling(rS), self.__getLeftSibling(lS) redistribute = True if (rS or lS) is not None: if rCnt == 2 or (rCnt == 1 and rSS != None and rSS.valcnt == 2): sib = rS elif lCnt == 2 or (lCnt == 1 and lSS != None and lSS.valcnt == 2): sib = lS elif lCnt == 1: sib, redistribute = lS, False elif rCnt == 1: sib, redistribute = rS, False if redistribute: # case 1: redistribute # left and right case if parent.valcnt == 1: if node == parent.getLink(0): parent_val, sib_val = parent.values[0], sib.values[0] child = sib.chooseChild(sib_val - 1) elif node == parent.getLink(1): parent_val, sib_val = parent.values[parent.valcnt - 1], sib.values[sib.valcnt - 1] child = sib.chooseChild(sib_val + 1) else: if sib == parent.getLink(1): # left if node == parent.getLink(0): parent_val, sib_val = parent.values[0], sib.values[0] child = sib.chooseChild(sib_val - 1) # right elif node == parent.getLink(2): parent_val, sib_val = parent.values[parent.valcnt - 1], sib.values[sib.valcnt - 1] child = sib.chooseChild(sib_val + 1) # middle elif sib == parent.getLink(2): parent_val, sib_val = parent.values[parent.valcnt - 1], sib.values[0] child = sib.chooseChild(sib_val - 1) elif sib == parent.getLink(0): parent_val, sib_val = parent.values[0], sib.values[sib.valcnt - 1] child = sib.chooseChild(sib_val + 1) node.insertValue(parent_val) parent.removeValue(parent_val) parent.insertValue(sib_val) sib.removeValue(sib_val) if not node.isLeafNode(): # if this is not a leaf node, redistribute the links also node.addLink(child) sib.removeLink(child) next_node = sib else: # case 2: merge if parent.valcnt == 1: parent_val = parent.values[0] else: if sib == parent.getLink(0): parent_val = parent.values[0] elif sib == parent.getLink(1): if sib == rS: parent_val = parent.values[0] if sib == lS: parent_val = parent.values[parent.valcnt - 1] child = node.getLink(0) sib.insertValue(parent_val) parent.removeValue(parent_val) parent.removeLink(node) if not node.isLeafNode(): sib.addLink(child) next_node = parent self.__fixNodeRemove(next_node, next_node.parent) else: # root node self.root = self.root.getLink(0) def __fixNodeInsert(self, node): if not node.isConsistent(): # conflict detected, try to resolve it if node.isLeafNode() and node is not self.root: # case for leaf node node.parent.insertValue(node.values[1]) node.parent.removeLink(node) # split the node node.parent.addLink(Node(node.values[0], node.parent)) node.parent.addLink(Node(node.values[node.valcnt - 1], node.parent)) self.__fixNodeInsert(node.parent) else: # case for internal node or root node if node is not self.root: node.parent.insertValue(node.values[1]) node.parent.removeLink(node) parent = node.parent else: self.root = Node(node.values[1]) parent = self.root # split the node leftNode, rightNode = Node(node.values[0], parent), Node(node.values[node.valcnt - 1], parent) parent.addLink(leftNode).addLink(rightNode) leftNode.addLink(node.getLink(0)).addLink(node.getLink(1)) rightNode.addLink(node.getLink(2)).addLink(node.getLink(3)) if node is not self.root: self.__fixNodeInsert(parent) # interface methods def contains(self, a): """ See if we have a given value in our tree """ node = self.findNode(a) return node if node.contains(a) else None def findNode(self, a): """ Find the node which contains the given value """ self.lastSearchDepth = 0 return self.__find(self.root, a) def findInorderSucc(self, node, a): """ Returns inorder successor of any node """ self.lastSearchDepth = 0 if node.isLeafNode(): return node new_node = node.chooseChild(a + 1) return self.__nextSucc(new_node) def insertValue(self, a): """ Inserts a new value to tree and keeps it balanced """ if self.root is None: self.root = Node(a) elif a is not None: node = self.findNode(a) res = node.contains(a) if res: return res # try to insert a new value into existing node node.insertValue(a) self.__fixNodeInsert(node) return self def insertList(self, xs): """ Insert a list of values into a tree """ if xs is not None and type(xs) is list: for item in xs: self.insertValue(item) def removeValue(self, a): """ Removes a value from the tree and keeps it balanced """ node = self.findNode(a) if not node or not node.contains(a): return None # swap the value we want to delete with its inorder successor (always leaf) succ = self.findInorderSucc(node, a) self.__swapValues(node, a, succ, succ.values[0]) # delete leaf node value succ.removeValue(a) # fix tree if needed self.__fixNodeRemove(succ) return self def removeList(self, xs): """ Deletes a list of values from a tree """ if xs is not None and type(xs) is list: for item in xs: self.removeValue(item)