skip list

import random


class SkipListNode(object):
    """Internal node for skip list."""

    def __init__(self, data, level):
        self.data = data
        self.level = level
        self.forwards = [None] * (level + 1)

    def get_forward(self, level):
        return self.forwards[level]

    def set_forward(self, level, node):
        self.forwards[level] = node

    def __repr__(self):
        return "<SkipListnode data: %s, level: %s>" % (self.data, self.level)

    def __str__(self):
        return str(self.data) + ":" + str(self.level)


class SkipList(object):
    """Skip list."""

    def __init__(self, size, prob):
        self.size = size
        self.prob = prob
        self.level = 0
        self.head = SkipListNode(None, self.size)
        random.seed()

    def _get_node_level(self):
        level = 0
        while random.randint(0, 10 ** 3) / float(10 ** 3) < self.prob:
            level += 1
        return min(level, self.size - 1)

    def _find_insert_position(self, data, level):
        """Find all see-able previous nodes."""
        seeables = []
        prev = self.head
        for lvl in reversed(list(range(level + 1))):
            node = prev.get_forward(lvl)
            while node and node.data < data:
                prev, node = node, node.get_forward(lvl)
            seeables.append(prev)
        return list(reversed(seeables))

    def find(self, data):
        prev = self.head
        for lvl in reversed(list(range(self.level))):
            node = prev.get_forward(lvl)
            while node and node.data < data:
                prev, node = node, node.get_forward(lvl)
            if node and node.data == data:
                return node
        return None

    def insert(self, data):
        level = self._get_node_level()
        self.level = max(self.level, level)
        node = SkipListNode(data, level)
        seeables = self._find_insert_position(data, level)
        for lvl, seeable in enumerate(seeables):
            node.set_forward(lvl, seeable.get_forward(lvl))
            seeable.set_forward(lvl, node)

    def remove(self, data):
        node = self.find(data)
        if node:
            seeables = self._find_insert_position(node.data, node.level)
            for lvl, seeable in enumerate(seeables):
                seeable.set_forward(lvl, node.get_forward(lvl))
            # update the skip list level
            for lvl in reversed(range(self.level + 1)):
                if self.head.get_forward(lvl):
                    self.level = lvl
                    break

    def print(self):
        for lvl in reversed(range(self.level + 1)):
            node = self.head.get_forward(lvl)
            while node:
                print(node, end=" ")
                node = node.get_forward(lvl)
            print()


if __name__ == "__main__":
    sl = SkipList(4, 0.5)
    for i in [10, 8, 3, 1, 2, 9, 7, 4, 6, 5]:
        sl.insert(i)
    print("After insertion:")
    sl.print()
    sl.remove(5)
    print("After removing 5:")
    sl.print()