How to randomly select a key based on its Integer value in a Map with respect to the other values in O(n) time?

Question

If we have a Map<T, Integer>, let's say the Integer value represents "how many" Ts there are. Thus, I want to uniformly select a T based on its Integer value. If the map contains Strings with "a"=4 and "b"=6, then I want it so that 40% of the time "a" is selected and 60% of the time "b" is selected.

Most importantly, I'd like this in O(n), n being two (not ten) in my previous example. I originally made an ArrayList containing the keys by how many values it had (and simply returning any random index), but this process is not only very slow, but completely counterintuitive for what the Map<T, Integer> represents.

Answer 1

Sorry for the delay, but I think I have a relatively elegant solution with O(n lg n) construction time and O(lg n) fetch-a-random-element time. Here goes.

---

WeightedProbMap: This class implements the random element generator. It is constructed based on an Iterable; see Test.java below.

import java.util.Random;
import java.util.SortedMap;
import java.util.TreeMap;

class WeightedProbMap<EltType>  {
    private SortedMap<Integer, EltType> elts = new TreeMap<Integer, EltType>();
    private Random rand = new Random();
    private int sum = 0;

    // assume: each weight is > 0; there is at least one element;
    //         elements should not be repeated
    // ensure: this.elts maps cumulative weights to elements;
    //         this.sum is the total weight
    public WeightedProbMap(Iterable<Pair<Integer, EltType>> weights) {
        for (Pair<Integer, EltType> e : weights) {
            this.elts.put(this.sum, e.second);
            this.sum += e.first;
        }
    }

    // assume: this was initialized properly (cf. constructor req)
    // ensure: return an EltType with relative probability proportional
    //         to its associated weight
    public EltType nextElt() {
        int index = this.rand.nextInt(this.sum) + 1;
        SortedMap<Integer, EltType> view = this.elts.headMap(index);
        return view.get(view.lastKey());
    }
}

---

Pair.java: Just a simple Pair class.

class Pair<X, Y> {
    public Pair(X x, Y y) {
        first = x;
        second = y;
    }

    X first;
    Y second;
}

---

Test.java: This is a very simple test harness for the WeightedProbMap (WPM) class. We build an ArrayList of elements with associated weights, use that to construct a WPM, and then get 10,000 samples from the WPM to see if elements appear with the expected frequency.

import java.util.ArrayList;

class Test {
    public static void main(String argc[]) {
        ArrayList<Pair<Integer, String> > elts = new ArrayList<Pair<Integer, String>>();
        elts.add(new Pair<Integer, String>(20, "Hello"));
        // elts.add(new Pair<Integer, String>(70, "World"));
        // elts.add(new Pair<Integer, String>(10, "Ohai"));

        WeightedProbMap<String> wpm = new WeightedProbMap<String>(elts);

        for (int i = 0; i < 10000; ++i) {
            System.out.println(wpm.nextElt());
        }
    }
}

---

Testing this:

1. Uncomment one or both of the elts.add(...) lines in Test.java.

2. Compile with:

   $ javac Pair.java WeightedProbMap.java Test.java

3. Run with (for example, in Unix):

   $ java Test | grep "Hello" | wc -l

This will give you the count for that particular execution.

---

Explanation:

constructor: The WeightedProbMap (WPM) class uses a java.util.SortedMap to map cumulative weights to elements. A graphical explanation:

The constructor takes weights...     ...and creates a mapping from the
      3 +---+                            number line:
        |   | 
  2 +---+   +---+ 2                   0      2         5      7
    |   |   |   |                     +------+---------+------+
    |   |   |   |                     |   X  |    Y    |   Z  |
  --+---+---+---+--                   +------+---------+------+
      X   Y   Z

nextElt(): A SortedMap stores its data by key order, which allows it to cheaply provide 'views' of subsets of the map. In particular, the line

SortedMap<Integer, EltType> view = this.elts.headMap(index)

returns a view of the original map (this.elts) with only the keys that are strictly smaller than index. This operation (headMap) is constant time: view takes O(1) time to construct, and if you were to change this.elts later on, the changes would be reflected in view as well.

Once we create the view of everything less than a random number, we now just have to find the greatest key in that subset. We do that with SortedMap.lastKey(), which, for a TreeMap, should take \Theta(lg n) time.