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 bette

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