Skip to content

Instantly share code, notes, and snippets.

# Plutor/splaytree.py Created Sep 4, 2012

Splay tree implementation in Python
 """Splay tree Logan Ingalls Sept 3, 2012 Note that I only implemented insert and find. Delete is trivial, and isn't the interesting part of splay trees. Some of these algorithms would be simpler, but I chose to do this without parent links from the tree nodes. Example output on my desktop computer: Building trees Done building Searched for 20 items 20000x in splay tree: 4.1 sec Searched for 20 items 20000x in binary tree: 11.1 sec """ from random import shuffle import time class BinaryTreeNode: """A node in a binary tree. Each node has a val attribute, as well as left and right descendents. No parent links because I'm a masochist.""" def __init__(self, val): self.val = val self.left = None self.right = None def __str__(self): return ("node[%s]" % self.val) class BinaryTree: """A basic binary tree.""" def __init__(self): self.root = None def insert(self, val, parent=None): """Insert a new value in the tree. Takes one argument (the value to insert). Recursive binary search.""" if (parent == None): parent = self.root if (parent == None): # root is null, make this the new root, done self.root = BinaryTreeNode(val) return elif (val < parent.val): if (parent.left == None): # insert to the left of the parent parent.left = BinaryTreeNode(val) return else: # search under the left self.insert(val, parent.left) else: if (parent.right == None): # insert to the right of the parent parent.right = BinaryTreeNode(val) return else: # search under the right self.insert(val, parent.right) def find(self, val, node=None): """Find if a value is in the tree. Takes one argument (the value to find). If the value is in the tree, returns the node object. Otherwise returns None.""" if (node == None): node = self.root if (node == None): # obviously it's not in an empty tree return None elif (val == node.val): return node elif (val < node.val): # Search left if (node.left != None): leftrv = self.find(val, node.left) if leftrv != None: return leftrv elif (val > node.val): if (node.right != None): rightrv = self.find(val, node.right) if rightrv != None: return rightrv return None class SplayTree(BinaryTree): """Implementation of a splay tree. Splay trees are self-organizing. They look identical to binary trees, but when nodes are found, they are moved towards the root (one or two levels closer). Still O(lg n), but a shallower search on average when not all values are searched for equally.""" def find(self, val, node=None, p=None, g=None, gg=None): """Find if a value is in the tree. Takes one argument (the value to find). If the value is in the tree, returns the node object. Otherwise returns None.""" if (node == None): node = self.root if (node == None): # obviously it's not in an empty tree return None elif (val == node.val): # If it's found, we need to move things around if (p != None): if (g == None): # Zig: swap node with its parent self.rotateup(node, p, g) elif ((g.left == p and p.left == node) or (g.right == p and p.right == node)): # Zig-zig: swap parent with grandparent self.rotateup(p, g, gg) # Then swap node with parent self.rotateup(node, p, gg) else: # Zig-zag: swap node with parent self.rotateup(node, p, g) # Then swap node with grandparent self.rotateup(node, g, gg) return node elif (val < node.val): # Search left if (node.left != None): leftrv = self.find(val, node.left, node, p, g) if leftrv != None: return leftrv elif (val > node.val): if (node.right != None): rightrv = self.find(val, node.right, node, p, g) if rightrv != None: return rightrv return None def rotateup(self, node, parent, gp=None): """Swap a node with its parent, keeping all child nodes (and grandparent node) in order.""" if node == parent.left: parent.left = node.right node.right = parent if (self.root == parent): self.root = node elif node == parent.right: parent.right = node.left node.left = parent if (self.root == parent): self.root = node else: print("This is impossible") if (gp != None): if (gp.right == parent): gp.right = node elif (gp.left == parent): gp.left = node def test_splay_tree(treesize=100000, iters=20000): """Just a simple test harness to demonstrate the speed of splay trees when a few items are searched for very frequently.""" # Build a binary tree and a splay tree print("Building trees") bintree = BinaryTree() spltree = SplayTree() x = [i for i in range(0, treesize)] shuffle(x) for n in x: bintree.insert(n) spltree.insert(n) print("Done building") searches = x[-20:] # Search the splay tree 1000 times t1 = time.time() for i in range(0, iters): for n in searches: node = spltree.find(n) if (node == None): print("ERROR: %d" % n) t2 = time.time() print("Searched for 20 items %dx in splay tree: %.1f sec" % (iters, t2-t1)) # Search the binary tree 1000 times t1 = time.time() for i in range(0, iters): for n in searches: node = bintree.find(n) if (node == None): print("ERROR: %d" % n) t2 = time.time() print("Searched for 20 items %dx in binary tree: %.1f sec" % (iters, t2-t1)) test_splay_tree()

### samueltcsantos commented May 2, 2013

 Very nice your code!

### wdevanny commented Oct 9, 2014

 This splay tree code is INCORRECT as of October 9th. Inserted and found nodes should be rotated all the way up to the root of the splay tree. In the code, found nodes are only rotated up two levels and inserted nodes are left untouched after being inserted as a leaf. Both of these behaviors are incorrect. See the first line of: http://en.wikipedia.org/wiki/Splay_tree#Splaying
to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.