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}