Constraints

• 1 <= len(operations) <= 200000
• Values are integers in the range [-10^9, 10^9]
• Snapshot ids and iterator ids used in the input are valid
• Target complexity: add/remove/contains in O(log n) expected, snapshot in O(1), iterator creation in O(log n), and next in O(1) amortized

Examples

Solution

def solution(operations):
    import random

    rng = random.Random(0)

    class Node:
        __slots__ = ('key', 'prio', 'left', 'right')
        def __init__(self, key, prio, left=None, right=None):
            self.key = key
            self.prio = prio
            self.left = left
            self.right = right

    def contains(root, x):
        while root is not None:
            if x < root.key:
                root = root.left
            elif x > root.key:
                root = root.right
            else:
                return True
        return False

    def merge(a, b):
        if a is None:
            return b
        if b is None:
            return a
        if a.prio < b.prio:
            return Node(a.key, a.prio, a.left, merge(a.right, b))
        else:
            return Node(b.key, b.prio, merge(a, b.left), b.right)

    def split(root, key):
        # returns (< key, >= key)
        if root is None:
            return None, None
        if key <= root.key:
            l, r_left = split(root.left, key)
            return l, Node(root.key, root.prio, r_left, root.right)
        else:
            l_right, r = split(root.right, key)
            return Node(root.key, root.prio, root.left, l_right), r

    def add(root, x):
        if contains(root, x):
            return root
        left, right = split(root, x)
        node = Node(x, rng.getrandbits(64))
        return merge(merge(left, node), right)

    def remove(root, x):
        # Remove x if present (safe even if absent)
        left, mid_right = split(root, x)
        mid, right = split(mid_right, x + 1)
        return merge(left, right)

    def push_left(node, stack):
        while node is not None:
            stack.append(node)
            node = node.left

    live_root = None
    snapshots = []  # list of roots
    iterators = []  # list of stacks for in-order traversal

    out = []
    for op in operations:
        cmd = op[0]
        if cmd == 'add':
            x = op[1]
            live_root = add(live_root, x)
        elif cmd == 'remove':
            x = op[1]
            live_root = remove(live_root, x)
        elif cmd == 'contains':
            x = op[1]
            out.append(contains(live_root, x))
        elif cmd == 'snapshot':
            snapshots.append(live_root)
            out.append(len(snapshots) - 1)
        elif cmd == 'iterator':
            sid = op[1]
            root = snapshots[sid] if 0 <= sid < len(snapshots) else None
            stack = []
            push_left(root, stack)
            iterators.append(stack)
            out.append(len(iterators) - 1)
        elif cmd == 'next':
            iid = op[1]
            if 0 <= iid < len(iterators):
                stack = iterators[iid]
                if stack:
                    node = stack.pop()
                    push_left(node.right, stack)
                    out.append(node.key)
                else:
                    out.append(None)
            else:
                out.append(None)
        else:
            # Unknown command: ignore or raise; per problem, not expected
            pass

    return out

Time complexity: add/remove/contains: O(log n) expected; snapshot: O(1); iterator creation: O(log n); next: O(1) amortized. Space complexity: O(log n) new nodes per add/remove operation, O(log n) per active iterator, plus O(number of snapshots) root references.

Hints

1. A snapshot can be just a pointer/reference to an immutable version of the set instead of a full copy.
2. If the set is stored in a persistent balanced BST, an iterator can keep a stack of the current in-order path and get amortized O(1) next() calls.