A light-weight implement of Splay Tree in Python (starter file)

splay_tree.py

# python3

__author__ = "github.com/duducheng"

'''A starter file for problem "set_range_sum" in Week 4 of 
[Data Structure](https://www.coursera.org/learn/data-structures/) course. '''


class SplayTree:
    def __init__(self):
        self.root = SplayEmpty()

    def __str__(self):
        field = ["key", "sum", "left", "right"]
        optinal = ",".join("{}={}".format(f, self.__dict__[f]) for f in field if f in self.__dict__)
        return "{}({})".format(self.__class__.__name__, optinal)

    def is_empty(self):
        '''To check if is `Empty`'''
        return isinstance(self, SplayEmpty)

    def not_empty(self):
        '''To check if is `Vertex`'''
        return isinstance(self, SplayVertex)

    def splay(self):
        pass

    def find(self, key):
        return self.root.find(key)

    def split(self, key):
        return self.root.split(key)

    @staticmethod
    def merge(node_left, node_right):
        if node_left.is_empty():
            return node_right
        if node_right.is_empty():
            return node_left
        while node_right.left.not_empty():
            node_right = node_right.left
        node_right.splay()
        node_right.left = node_left
        return node_right

    def insert(self, x):
        '''operation +'''
        self.root = self.root.insert(x)

    def erase(self, x):
        '''operation -'''
        self.root = self.root.erase(x)

    def search(self, x):
        '''operation ?'''
        return self.root.search(x)

    def sum_from_to(self, fr, to):
        '''operation s'''
        return self.root.sum_from_to(fr, to)


class SplayEmpty(SplayTree):
    '''Empty means `None` in the Tree.'''
    sum = 0
    parent = None
    left = None
    right = None

    def __init__(self):
        pass

    def find(self, key):
        return self

    def split(self, key):
        return self, self

    def insert(self, key):
        return SplayVertex(key)

    def erase(self, key):
        return self

    def search(self, key):
        return False

    def sum_from_to(self, fr, to):
        return 0


class SplayVertex(SplayTree):
    def __init__(self, key, left=SplayEmpty(), right=SplayEmpty()):
        self.parent = self._left = self._right = SplayEmpty()
        self.key = self._sum = key
        self.left = left
        self.right = right

    @property
    def sum(self):
        return self._sum

    @property
    def left(self):
        return self._left

    @property
    def right(self):
        return self._right

    @left.setter
    def left(self, node):
        '''When change the child, do update with its' sum and parent.'''
        self._left = node
        self._sum = self.key + node.sum + self.right.sum
        node.parent = self

    @right.setter
    def right(self, node):
        '''When change the child, do update with its' sum and parent.'''
        self._right = node
        self._sum = self.key + node.sum + self.left.sum
        node.parent = self

    @property
    def root(self):
        if self.parent.is_empty():
            return self
        else:
            return self.parent.root

    def _small_rotation(self):
        parent = self.parent
        if parent.is_empty():
            return
        grandparent = self.parent.parent
        if parent.left == self:
            m = self.right
            parent.left = m
            self.right = parent
        else:
            m = self.left
            parent.right = m
            self.left = parent
        if grandparent.not_empty():
            if grandparent.left == parent:
                grandparent.left = self
            else:
                grandparent.right = self
        else:
            self.parent = SplayEmpty()

    def _big_rotation(self):
        if (self.parent.left == self and self.parent.parent.left == self.parent) or \
                (self.parent.right == self and self.parent.parent.right == self.parent):
            # `Zig-zig`
            self.parent._small_rotation()
            self._small_rotation()
        else:
            # `Zig-zag`
            self._small_rotation()
            self._small_rotation()

    def splay(self):
        '''Makes splay of the given vertex and makes it the new root. '''
        while self.parent.not_empty():
            if self.parent.parent.is_empty():
                self._small_rotation()  # `Zig`
                break
            self._big_rotation()

    def find(self, key):
        '''Searches for the given key in the tree with the given root and
        calls splay for the deepest visited node after that.

        Return the found result.
            - If found, result is a pointer to the node with the given key.
            - Otherwise, result is a pointer to the node with the smallest
            bigger key (next value in the order).
            - If the key is bigger than all keys in the tree,
            then result is Empty.
        '''
        v = self.root
        last_access = self.root
        next_access = SplayEmpty()
        while v.not_empty():
            if v.key >= key and (next_access.is_empty() or v.key < next_access.key):
                next_access = v
            last_access = v
            if v.key == key:
                break
            if v.key < key:
                v = v.right
            else:
                v = v.left
        last_access.splay()
        return next_access

    def split(self, key):
        result = self.find(key)
        if result.is_empty():
            return self.root, SplayEmpty()
        result.splay()
        node_right = result
        node_left = node_right.left
        node_right.left = SplayEmpty()
        if node_left.not_empty():
            node_left.parent = SplayEmpty()
        return node_left, node_right

    def insert(self, key):
        node_left, node_right = self.split(key)
        if node_right.is_empty() or node_right.key != key:
            new_vertex = __class__(key)
        else:
            new_vertex = SplayEmpty()
        return self.merge(self.merge(node_left, new_vertex), node_right)

    def erase(self, key):
        (node_left, node_right) = self.split(key)
        # TODO:Implement erase yourself

        return self

    def search(self, key):
        # TODO:Implement `search` yourself

        return False

    def sum_from_to(self, fr, to):
        node_left, node_middle = self.split(fr)
        node_middle, node_right = node_middle.split(to + 1)
        # TODO: Implement `sum_from_to` yourself

        ans = 0
        return ans


if __name__ == "__main__":
    from sys import stdin

    MODULO = 1000000001
    n = int(stdin.readline())
    splay_tree = SplayTree()
    last_sum_result = 0
    for i in range(n):
        line = stdin.readline().split()
        if line[0] == '+':
            x = int(line[1])
            splay_tree.insert((x + last_sum_result) % MODULO)
        elif line[0] == '-':
            x = int(line[1])
            splay_tree.erase((x + last_sum_result) % MODULO)
        elif line[0] == '?':
            x = int(line[1])
            print('Found' if splay_tree.search((x + last_sum_result) % MODULO) else 'Not found')
        elif line[0] == 's':
            l = int(line[1])
            r = int(line[2])
            res = splay_tree.sum_from_to((l + last_sum_result) % MODULO, (r + last_sum_result) % MODULO)
            print(res)
            last_sum_result = res % MODULO