Python data structure for efficient add, remove, and random.choice

I’m looking for a built-in Python data structure that can add a new element, remove an existing element, and choose a random element, all in better than O(n) time.

I was hoping that set could do this, but AFAIK, the only way to choose a random element from a Python set is random.choice(list(my_set)), which takes O(n) time.

I would greatly prefer a solution that’s built into Python, since I require efficiency and easy deployment. Unfortunately, Python does not seem to have built-in tree data types.

Best answer

Python does not have a built-in data structure which meets all 3 of your requirements.

That said, it’s fairly trivial to implement a tree yourself.


Another option would be to combine a dictionary with a list to create what is effectively a set that also maintains a list of its items:

import random

class ListDict(object):
    def __init__(self):
        self.item_to_position = {}
        self.items = []

    def add_item(self, item):
        if item in self.item_to_position:
            return
        self.items.append(item)
        self.item_to_position[item] = len(self.items)-1

    def remove_item(self, item):
        position = self.item_to_position.pop(item)
        last_item = self.items.pop()
        if position != len(self.items):
            self.items[position] = last_item
            self.item_to_position[last_item] = position

    def choose_random_item(self):
        return random.choice(self.items)

Since the only operations done on the list are .pop() and .append(), they shouldn’t take more than constant time (in most Python implementations, at least).