How to get a random value in a very large python dictionary

Question

Given a python dict with multiple million entries, what is the most efficient way to get and remove a random (k,v) pair from it?

The dict is constantly growing and the random remove function is called very often.

The most cited solution for python2 random_key = random.choice(the_dict.keys()) is way too slow as a list of all the keys is created first. With lots of elements in the dict, this solution does not work.

Another proposed solution is the_dict.popitem() , but this does not return a real random object, but depends on the internal ordering of the dict.

The third solution that is also way to slow is an iterator:

 it = the_dict.iterkeys()

 for i in range (random.randint(0, len(the_dict)-1)):
     next(it)
 random_key = next(it)

Next to the remove_random() , sometimes a the_dict.pop(x) is required for a specific key. Therefore, a simple list based secondary index does not work.

Can this problem be efficiently be solved with a dict?

Answer 1

A solution is to use a bidirectional map each key to an integer, to allow for randomly selecting a key by using random.randrange(0,N) to select from a range of ints which are bidirectionally mapped to the keys, where N is the number of keys.

Adding a new key simply assigns it the next higher int. Deleting a key reassigns the int for that key to the key that was assigned the previously highest int before deleting the key-value pair. Python code provided for clarity.

Python code:

def create(D): # O(len(D))
    # Create the bidirectional maps from the dictionary, D
    keys = D.keys()
    ints = range(len(keys)
    int_to_key = dict(zip(keys, ints)) 
    key_to_int = dict(zip(ints, keys))
    return (int_to_key, key_to_int)

def add(D, int_to_key, key_to_int, key, value): # O(1)
    # Add key-value pair (no extra work needed for simply changing the value)
    new_int = len(D)
    D[key] = value
    int_to_key[new_int] = key
    key_to_int[key] = new_int

def remove(D, int_to_key, key_to_int, key): # O(1)
    # Update the bidirectional maps then remove the key-value pair

    # Get the two ints and keys.
    key_int = key_to_int[key]
    swap_int = len(D) - 1 # Should be the highest int
    swap_key = int_to_key[swap_int]

    # Update the bidirectional maps so that key now has the highest int
    key_to_int[key], key_to_int[swap_key] = swap_int, key_int
    int_to_key[key_int], int_to_key[swap_int] = swap_key, key

    # Remove elements from dictionaries
    D.remove(key)
    key_to_int.remove(key)
    int_to_key.remove(key)

def random_key(D, int_to_key): # O(1)
    # Select a random key from the dictionary using the int_to_key map
    return int_to_key[random.randrange(0, len(D))]

def remove_random(D, int_to_key, key_to_int): # O(1)
    # Randomly remove a key from the dictionary via the bidirectional maps
    key = random_key(D, int_to_key)
    remove(D, int_to_key, key_to_int, key)

Note: Adding/removing keys from D without using the appropriate above functions will break the bidirectional map. This means it's best to implement this as a class.

Answer 2

No, as you've discovered, this can't be done efficiently with a plain dict. See this issue for some explanations about why implementing random.choice for sets is hard; the same arguments apply to dictionaries.

But it's possible to create a dict-like data structure that does support efficient random selection. Here's a recipe for such an object, based in part on this question and its responses. It's only a starting-point, but it supports most of the existing dict methods, many of which are conveniently filled in by the MutableMapping ABC . Depending on your needs, you may need to flesh it out a bit: for example, to be able to create a RandomChoiceDict directly from a regular dict, or to add a meaningful __repr__ , etc.

Essentially, you need to maintain three structures: a list of keys, a list of corresponding values, and a dict that maps keys back to indices (the inverse of the keys list). The basic __getitem__ , __setitem__ and __delitem__ operations can be simply implemented in terms of those structures, and if __len__ and __iter__ are specified, the abstract base class takes care of most of the rest.

from collections import MutableMapping
import random

class RandomChoiceDict(MutableMapping):
    """
    Dictionary-like object allowing efficient random selection.

    """
    def __init__(self):
        # Add code to initialize from existing dictionaries.
        self._keys = []
        self._values = []
        self._key_to_index = {}

    def __getitem__(self, key):
        return self._values[self._key_to_index[key]]

    def __setitem__(self, key, value):
        try:
            index = self._key_to_index[key]
        except KeyError:
            # Key doesn't exist; add a new one.
            index = len(self._keys)
            self._key_to_index[key] = index
            self._keys.append(key)
            self._values.append(value)
        else:
            # Key already exists; overwrite the value.
            self._values[index] = value

    def __delitem__(self, key):
        index = self._key_to_index.pop(key)
        # Remove *last* indexed element, then put
        # it back at position 'index' (overwriting the
        # one we're actually removing) if necessary.
        key, value = self._keys.pop(), self._values.pop()
        if index != len(self._key_to_index):
            self._keys[index] = key
            self._values[index] = value
            self._key_to_index[key] = index

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

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

    def random_key(self):
        """Return a randomly chosen key."""
        if not self:
            raise KeyError("Empty collection")
        index = random.randrange(len(self))
        return self._keys[index]

    def popitem_random(self):
        key = self.random_key()
        value = self.pop(key)
        return key, value

Example usage:

>>> d = RandomChoiceDict()
>>> for x in range(10**6):  # populate with some values
...     d[x] = x**2
... 
>>> d.popitem_random()  # remove and return random item
(132545, 17568177025)
>>> 132545 in d
False
>>> d.popitem_random()
(954424, 910925171776)