algorithm-in-python/dataStructure/splayTree.py

''' mbinary
#########################################################################
# File : splayTree.py
# Author: mbinary
# Mail: zhuheqin1@gmail.com
# Blog: https://mbinary.xyz
# Github: https://github.com/mbinary
# Created Time: 2018-05-19  23:06
# Description:
#########################################################################
'''

from collections import deque, Iterable
# use isinstance(obj,Iterable)  to judge if an obj is iterable


class node:
    def __init__(self, val=None, left=None, right=None, parent=None):
        self.val = val
        if val:
            self.freq = 1
        else:
            self.freq = 0
        self.left = left
        self.right = right
        self.parent = parent

    def getChild(self, s=0):
        if isinstance(s, int):
            s = [s]
        last = self
        for i in s:
            if not last:
                return None
            if i == 0:
                last = last.left
            else:
                last = last.right
        return last

    def setChild(self, child, s=0):
        if isinstance(s, Iterable):
            i = s[0]
            del s[0]
            if i == 0:
                self.left.setChild(child, s)
            else:
                self.right.setChild(child, s)
        elif s:
            self.right = child
        else:
            self.left = child


class splayTree:
    def __init__(self, s=[]):
        s = list(s)
        self.root = None
        s = sorted(s, reverse=True)
        for i in s:
            self.insert(self.root, i)

    def insert(self, k):
        if not self.root:
            self.root = node(k)
        else:
            self._insert(self.root, k)

    def _insert(self, root, k):
        if root.val == k:
            root.freq += 1
        elif root.val < k:
            if not root.right:
                root.right = node(k)
                root.right.parent = root
            else:
                self._insert(root.right, k)
        else:
            if not root.left:
                root.left = node(k)
                root.left.parent = root
            else:
                self._insert(root.left, k)

    def _zigzagRotate(self, i, j, root, parent, grand):
        parent.setChild(root.getChild(i), j)
        root.setChild(parent, i)
        grand.setChild(root.getChild(j), i)
        root.setChild(grand, j)
        if root.parent:
            root.parent = grand.parent
        parent.parent = root
        grand.parent = root

    def _lineRotate(self, i, root, parent, grand):
        grand.setChild(parent.getChild(i ^ 1), i)
        parent.setChild(grand, i ^ 1)
        parent.setChild(root.getChild(i ^ 1), i)
        root.setChild(parent, i ^ 1)
        if root.parent:
            root.parent = grand.parent
        parent.parent = root
        grand.parent = parent

    def _rotate(self, root):
        if root == self.root:
            return
        if root.parent == self.root:
            for i in range(2):
                if root.parent.getChild(i) == root:
                    root.parent.parent = root
                    root.parent.setChild(root.getChild(i ^ 1), i)
                    root.parent = None
                    root.setChild(self.root, i ^ 1)
                    self.root = root
        else:
            grand = root.parent.parent
            parent = root.parent
            if grand == self.root:
                self.root = root
                root.parent = None
            else:
                for i in range(2):
                    if grand.parent.getChild(i) == grand:
                        grand.parent.setChild(root, i)
            for i in range(2):
                for j in range(2):
                    if i != j and grand.getChild([i, j]) == root:
                        self._zigzagRotate(i, j, root, parent, grand)
                    elif i == j and grand.getChild([i, i]) == root:
                        self._lineRotate(i, root, parent, grand)
        self._rotate(root)

    def _find(self, root, k):
        if not root:
            return 0
        if root.val > k:
            return self._find(root.left, k)
        elif root.val < k:
            return self._find(root.right, k)
        else:
            self._rotate(root)
            return root.freq

    def _maxmin(self, root, i=0):
        if not root:
            return None
        if root.getChild(i):
            return self._maxmin(root.getChild(i))
        return root

    def Max(self):
        return self._maxmin(self.root, 1)

    def Min(self):
        return self._maxmin(self.root, 0)

    def remove(self, k):
        tmp = self.find(k)
        if not tmp:
            raise ValueError
        else:
            if self.root.left:
                r = self.root.right
                self.root = self.root.left
                self.root.parent = None
                Max = self.Max()
                Max.right = r
                if r:
                    r.parent = Max
            else:
                self.root = self.root.right

    def find(self, k):
        return self._find(self.root, k)

    def levelTraverse(self):
        q = deque()
        q.append((self.root, 0))
        rst = []
        while q:
            tmp, n = q.popleft()
            rst.append(tmp)
            if tmp.left:
                q.append((tmp.left, n+1))
            if tmp.right:
                q.append((tmp.right, n+1))
        return rst

    def display(self):
        data = self.levelTraverse()
        for i in data:
            print(i.val, end=' ')
        print('')


if __name__ == '__main__':
    a = splayTree()
    a.insert(5)
    a.insert(1)
    a.insert(4)
    a.insert(3)
    a.insert(2)
    a.insert(7)
    a.insert(8)
    a.insert(2)
    print('initial:5,1,4,2,7,8,2')
    a.display()
    tmp = a.find(2)
    print("after find(2):%d" % tmp)
    a.display()
    print("remove(4)")
    a.remove(4)
    a.display()