How to randomly select a key based on its Integer

问题:

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.

回答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.

回答2:

To do this, you have to cache the relative frequency of each value T. This gives you your O(n) probability-distribution for the price of an O(n) insertion-cost (you have to update the relative frequency of every T upon every insertion).

回答3:

If you can store the total sum, that's quite easily done:

Just store the pairs (T, int) as a class or whatever in an ordinary array and then go over it:

int val = Random.nextInt(total);
for (Pair p : pairs) {
    val -= p.val;
    if (val < 0) return p;
}

Can't get much faster considering that looping through an ArrayList is the most efficient way to iterate through n values and you can obviously not do better than O(n). The only overhead is the nextInt() and you need that (or something similar) as well in every solution. Depending on how you organize the ArrayList (sorted or not) other operations get cheaper/more expensive, but it's unimportant for that particular action

Edit: Although thinking about it the "you obviously need O(n)" isn't true. If you change the values in the array rarely and can allow a expensive preparation and memory isn't a problem you can do better by storing a HashMap. If you've got for example a distribution: T0: 2 T1: 3 T2: 1

You could insert (0, T0), (1, T0), (2, T1),.,(4, T1), (5, T2) in the hashmap.

Edit2: Or see phooji's approach which should be feasible for larger sets of data.

回答4:

Build an inverse map, Map<Integer,T>so that every key is the sum of all the weights processed so far.

For example if you have this map:

T1 -> 10
T2 -> 8
T3 -> 3

This inverse map is:

10 -> T1
18 -> T2
21 -> T3

(For better performance, you can arrange your weights in descending order first.)

Then generate an evenly distributed random number between 0 and the sum of all weights, and perform a binary search for this number in the key set of the inverse map.

回答5:

Using an arraylist would actually be even faster than using a Map, because you can do it in O(1).

class RandVal<T> {

    List<T> list = new ArrayList<T>();
    Random rand = new Random();

    public T randomValue() {
        int next = rand.nextInt(list.size());
        return list.get(next);
    }

}

The only way this is a bad thing is if order matters (A A B B A B vs A B B A B A or something) but it's obvious it doesn't because you're using a Map which has no ordering...

回答6:

OP here.

I came up with an elegant solution! For any misunderstandings: My original idea of storing all the keys by how many values in an ArrayList was completely disregarding the point of using a Map to store "instances of the Key using Integers"; any similar solutions are counterproductive! Assuming the Map is unordered, here is my solution:

public T randomPick(Random r) {

        int randomValue = r.nextInt(size());
        int currentSum = 0;
        T lastElement = null;

        for (T t : map.keySet()){
            if (randomValue < currentSum + map.get(t)){
                return t;
            }
            currentSum+= map.get(t);
            lastElement = t;
        }
        return lastElement;
    }

It compares the random value with a current sum + the current element's value. If it is less than that, we return the current key. Else, keep going and add that value to the sum. If it is the case such that the random value is never less than any of the values, we return the lastElement.

Hope this clears it up.