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:
- Uncomment one or both of the
elts.add(...) lines in Test.java.
Compile with:
$ javac Pair.java WeightedProbMap.java Test.java
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.
