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

Question

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.

Answer 1

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(), .append(), and index retrieval and assignment, they shouldn't take more than constant time (in most Python implementations, at least).

You can extend the above definition with extra methods to support other useful operations, like len, in, and iteration:

class ListDict(object):
    ... # methods from above

    def __contains__(self, item):
        return item in self.item_to_position

    def __iter__(self):
        return iter(self.items)

    def __len__(self):
        return len(self.items)