skiplist.py

import math
import random
  
class Node:
    __slots__ = ["key", "forward", "backward", "next"]
    
    def __init__(self, key, level):
        self.key = key
        self.forward = [None] * (level + 1)
        self.backward = [None] * (level + 1)
        
        self.next = None
  
class SkipList:
    def __init__(self, depth, p=0.25):
        self.depth = depth
        self.p = p

        # Skip list
        self.header = Node(-math.inf, depth)
        self.tail = Node(math.inf, depth)
        self.header.forward = [self.tail] * (self.depth + 1)
        self.tail.backward = [self.header] * (self.depth + 1)
        self.level = 0
        
        # Linked list
        self.first = None
        self.last = None
      
    def random_level(self):
        level = 0
        while random.random() < self.p:
            level += 1
            if level == self.depth:
                break
        return level
    
    def __getitem__(self, i):
        node = self.header
        for _ in range(i + 1):
            node = node.forward[0]
        return node
    
    def search(self, key):
        update = [self.header] * (self.depth + 1)
        current = self.header
  
        for i in range(self.level, -1, -1):
            while (current.forward[i] and
                   current.forward[i].key < key):
                current = current.forward[i]
            update[i] = current
        return update
        
    def popleft(self):
        node = self.first
        self.delete_element(node)
        self.first = node.next
        return node
  
    def push(self, key):
        # Update skip list
        update = self.search(key)
  
        random_level = self.random_level()

        if random_level > self.level:
            self.level = random_level

        node = Node(key, random_level)

        for i in range(random_level + 1):
            node.forward[i] = update[i].forward[i]
            node.backward[i] = update[i]
            update[i].forward[i].backward[i] = node
            update[i].forward[i] = node
            
        # Update linked list
        if self.last is not None:
            self.last.next = node
        else:
            self.first = node
        self.last = node
        
        return node
                
    def delete_element(self, node):
        for i, prev_node in enumerate(node.backward):
            prev_node.forward[i] = node.forward[i]
        for i, next_node in enumerate(node.forward):
            next_node.backward[i] = node.backward[i]

        while (self.level > 0 and
               self.header.forward[self.level] is self.tail):
            self.level -= 1