Search code examples
algorithmshuffle

Choosing k out of n


I want to choose k elements uniformly at random out of a possible n without choosing the same number twice. There are two trivial approaches to this.

  1. Make a list of all n possibilities. Shuffle them (you don't need to shuffle all n numbers just k of them by performing the first k steps of Fisher Yates). Choose the first k. This approach takes O(k) time (assuming allocating an array of size n takes O(1) time) and O(n) space. This is a problem if k is very small relative to n.
  2. Store a set of seen elements. Choose a number at random from [0, n-1]. While the element is in the set then choose a new number. This approach takes O(k) space. The run-time is a little more complicated to analyze. If k = theta(n) then the run-time is O(k*lg(k))=O(n*lg(n)) because it is the coupon collector's problem. If k is small relative to n then it takes slightly more than O(k) because of the probability (albeit low) of choosing the same number twice. This is better than the above solution in terms of space but worse in terms of run-time.

My question:

is there an O(k) time, O(k) space algorithm for all k and n?


Solution

  • With an O(1) hash table, the partial Fisher-Yates method can be made to run in O(k) time and space. The trick is simply to store only the changed elements of the array in the hash table.

    Here's a simple example in Java:

    public static int[] getRandomSelection (int k, int n, Random rng) {
        if (k > n) throw new IllegalArgumentException(
            "Cannot choose " + k + " elements out of " + n + "."
        );
    
        HashMap<Integer, Integer> hash = new HashMap<Integer, Integer>(2*k);
        int[] output = new int[k];
    
        for (int i = 0; i < k; i++) {
            int j = i + rng.nextInt(n - i);
            output[i] = (hash.containsKey(j) ? hash.remove(j) : j);
            if (j > i) hash.put(j, (hash.containsKey(i) ? hash.remove(i) : i));
        }
        return output;
    }
    

    This code allocates a HashMap of 2×k buckets to store the modified elements (which should be enough to ensure that the hash table is never rehashed), and just runs a partial Fisher-Yates shuffle on it.

    Here's a quick test on Ideone; it picks two elements out of three 30,000 times, and counts the number of times each pair of elements gets chosen. For an unbiased shuffle, each ordered pair should appear approximately 5,000 (±100 or so) times, except for the impossible cases where both elements would be equal.