Leetcode 1337 Solution

class ListNode:
    def __init__(self, val, prev=None, next=None, down=None):
        self.val = val
        self.next = next
        self.prev = prev
        self.down = down


class Skiplist:
    def __init__(self):
        self.heads = []
        self.tails = []

        for i in range(30):
            self.add_list()

    def print(self):
        for i in range(len(self.heads) - 1, -1, -1):
            head = self.heads[i]
            tail = self.tails[i]

            node = head.next
            values = []
            while node != tail:
                values.append(str(node.val))
                node = node.next
            print("List {}: {}".format(i, " -> ".join(values) if values else None))

    def add_list(self):
        head = ListNode(-1)
        tail = ListNode(float("inf"))
        head.next = tail
        tail.prev = head
        head.down = self.heads[-1] if self.heads else None

        self.heads.append(head)
        self.tails.append(tail)

    def search_nodes(self, target):
        ans = []

        node = self.heads[-1]
        for i in range(len(self.heads) - 1, -1, -1):
            tail = self.tails[i]

            while node != tail:
                if node.val <= target < node.next.val:
                    ans.append(node)
                    node = node.down
                    break
                node = node.next

        return list(reversed(ans))

    def search(self, target: int) -> bool:
        nodes = self.search_nodes(target)
        return nodes[0].val == target

    def add(self, num: int) -> None:
        nodes = self.search_nodes(num)

        last_node = None
        for i in range(len(nodes)):
            node = nodes[i]
            new_node = ListNode(num, node, node.next)
            node.next.prev = new_node
            node.next = new_node
            new_node.down = last_node
            last_node = new_node

            if random.random() < 0.5:
                break

    def erase(self, num: int) -> bool:
        nodes = self.search_nodes(num)

        if nodes[0].val != num:
            return False

        for i in range(len(nodes)):
            node = nodes[i]
            if node.val != num:
                break
            node.prev.next = node.next
            node.next.prev = node.prev

        return True