001package squidpony.squidmath;
002
003import squidpony.annotation.Beta;
004
005import java.io.Serializable;
006import java.util.HashMap;
007import java.util.Set;
008
009/**
010 * A generic method of holding a probability table to determine weighted random
011 * outcomes.
012 *
013 * The weights do not need to add up to any particular value, they will be
014 * normalized when choosing a random entry.
015 *
016 * @author Eben Howard - http://squidpony.com - howard@squidpony.com
017 *
018 * @param <T> The type of object to be held in the table
019 */
020@Beta
021public class ProbabilityTable<T> implements Serializable {
022    private static final long serialVersionUID = -1307656083434154736L;
023    private final HashMap<T, Integer> table = new HashMap<>();
024    private RNG rng;
025    private int total = 0;
026
027    /**
028     * Creates a new probability table.
029     */
030    public ProbabilityTable() {
031        this(new StatefulRNG());
032    }
033
034    /**
035     * Creates a new probability table with the provided source of randomness
036     * used.
037     *
038     * @param rng the source of randomness
039     */
040    public ProbabilityTable(RNG rng) {
041        this.rng = rng;
042    }
043
044    /**
045     * Returns an object randomly based on assigned weights.
046     *
047     * Returns null if no elements have been put in the table.
048     *
049     * @return the chosen object or null
050     */
051    public T random() {
052        if (table.isEmpty()) {
053            return null;
054        }
055        int index = rng.nextInt(total);
056        for (T t : table.keySet()) {
057            index -= table.get(t);
058            if (index < 0) {
059                return t;
060            }
061        }
062        return null;//something went wrong, shouldn't have been able to get all the way through without finding an item
063    }
064
065    /**
066     * Adds the given item to the table.
067     *
068     * Weight must be greater than 0.
069     *
070     * @param item the object to be added
071     * @param weight the weight to be given to the added object
072     */
073    public void add(T item, int weight) {
074        Integer i = table.get(item);
075        if (i == null) {
076            i = weight;
077        } else {
078            i += weight;
079        }
080        table.put(item, i);
081        total += weight;
082    }
083
084    /**
085     * Returns the weight of the item if the item is in the table. Returns zero
086     * of the item is not in the table.
087     *
088     * @param item the item searched for
089     * @return the weight of the item, or zero
090     */
091    public int weight(T item) {
092        Integer i = table.get(item);
093        return i == null ? 0 : i;
094    }
095
096    /**
097     * Provides a set of the items in this table, without reference to their
098     * weight.
099     *
100     * @return a set of all items stored
101     */
102    public Set<T> items() {
103        return table.keySet();
104    }
105}