third-party/net.sf.jung2/src/main/java/edu/uci/ics/jung/algorithms/util/WeightedChoice.java

   1 /**
   2  * Copyright (c) 2009, the JUNG Project and the Regents of the University
   3  * of California
   4  * All rights reserved.
   5  *
   6  * This software is open-source under the BSD license; see either
   7  * "license.txt" or
   8  * http://jung.sourceforge.net/license.txt for a description.
   9  * Created on Jan 8, 2009
  10  *
  11  */
  12 package edu.uci.ics.jung.algorithms.util;
  13
  14 import java.util.ArrayList;
  15 import java.util.LinkedList;
  16 import java.util.List;
  17 import java.util.Map;
  18 import java.util.Queue;
  19 import java.util.Random;
  20
  21 /**
  22  * Selects items according to their probability in an arbitrary probability
  23  * distribution.  The distribution is specified by a {@code Map} from
  24  * items (of type {@code T}) to weights of type {@code Number}, supplied
  25  * to the constructor; these weights are normalized internally to act as
  26  * probabilities.
  27  *
  28  * <p>This implementation selects items in O(1) time, and requires O(n) space.
  29  *
  30  * @author Joshua O'Madadhain
  31  */
  32 public class WeightedChoice<T>
  33 {
  34         private List<ItemPair> item_pairs;
  35         private Random random;
  36
  37         /**
  38          * The default minimum value that is treated as a valid probability
  39          * (as opposed to rounding error from floating-point operations).
  40          */
  41         public static final double DEFAULT_THRESHOLD = 0.00000000001;
  42
  43         /**
  44          * Equivalent to {@code this(item_weights, new Random(), DEFAULT_THRESHOLD)}.
  45          * @param item_weights
  46          */
  47         public WeightedChoice(Map<T, ? extends Number> item_weights)
  48         {
  49                 this(item_weights, new Random(), DEFAULT_THRESHOLD);
  50         }
  51
  52         /**
  53          * Equivalent to {@code this(item_weights, new Random(), threshold)}.
  54          */
  55         public WeightedChoice(Map<T, ? extends Number> item_weights, double threshold)
  56         {
  57                 this(item_weights, new Random(), threshold);
  58         }
  59
  60         /**
  61          * Equivalent to {@code this(item_weights, random, DEFAULT_THRESHOLD)}.
  62          */
  63         public WeightedChoice(Map<T, ? extends Number> item_weights, Random random)
  64         {
  65                 this(item_weights, random, DEFAULT_THRESHOLD);
  66         }
  67
  68         /**
  69          * Creates an instance with the specified mapping from items to weights,
  70          * random number generator, and threshold value.
  71          *
  72          * <p>The mapping defines the weight for each item to be selected; this
  73          * will be proportional to the probability of its selection.
  74          * <p>The random number generator specifies the mechanism which will be
  75          * used to provide uniform integer and double values.
  76          * <p>The threshold indicates default minimum value that is treated as a valid
  77          * probability (as opposed to rounding error from floating-point operations).
  78          */
  79         public WeightedChoice(Map<T, ? extends Number> item_weights, Random random,
  80                         double threshold)
  81         {
  82                 if (item_weights.isEmpty())
  83                         throw new IllegalArgumentException("Item weights must be non-empty");
  84
  85                 int item_count = item_weights.size();
  86                 item_pairs = new ArrayList<ItemPair>(item_count);
  87
  88                 double sum = 0;
  89                 for (Map.Entry<T, ? extends Number> entry : item_weights.entrySet())
  90                 {
  91                         double value = entry.getValue().doubleValue();
  92                         if (value <= 0)
  93                                 throw new IllegalArgumentException("Weights must be > 0");
  94                         sum += value;
  95                 }
  96         double bucket_weight = 1.0 / item_weights.size();
  97
  98                 Queue<ItemPair> light_weights = new LinkedList<ItemPair>();
  99                 Queue<ItemPair> heavy_weights = new LinkedList<ItemPair>();
 100                 for (Map.Entry<T, ? extends Number> entry : item_weights.entrySet())
 101                 {
 102                         double value = entry.getValue().doubleValue() / sum;
 103                         enqueueItem(entry.getKey(), value, bucket_weight, light_weights, heavy_weights);
 104                 }
 105
 106                 // repeat until both queues empty
 107                 while (!heavy_weights.isEmpty() || !light_weights.isEmpty())
 108                 {
 109                         ItemPair heavy_item = heavy_weights.poll();
 110                         ItemPair light_item = light_weights.poll();
 111                         double light_weight = 0;
 112                         T light = null;
 113                         T heavy = null;
 114                         if (light_item != null)
 115                         {
 116                                 light_weight = light_item.weight;
 117                                 light = light_item.light;
 118                         }
 119                         if (heavy_item != null)
 120                         {
 121                                 heavy = heavy_item.heavy;
 122                                 // put the 'left over' weight from the heavy item--what wasn't
 123                                 // needed to make up the difference between the light weight and
 124                                 // 1/n--back in the appropriate queue
 125                                 double new_weight = heavy_item.weight - (bucket_weight - light_weight);
 126                                 if (new_weight > threshold)
 127                                         enqueueItem(heavy, new_weight, bucket_weight, light_weights, heavy_weights);
 128                         }
 129                         light_weight *= item_count;
 130
 131                         item_pairs.add(new ItemPair(light, heavy, light_weight));
 132                 }
 133
 134                 this.random = random;
 135         }
 136
 137         /**
 138          * Adds key/value to the appropriate queue.  Keys with values less than
 139          * the threshold get added to {@code light_weights}, all others get added
 140          * to {@code heavy_weights}.
 141          */
 142         private void enqueueItem(T key, double value, double threshold,
 143                         Queue<ItemPair> light_weights, Queue<ItemPair> heavy_weights)
 144         {
 145                 if (value < threshold)
 146                         light_weights.offer(new ItemPair(key, null, value));
 147                 else
 148                         heavy_weights.offer(new ItemPair(null, key, value));
 149         }
 150
 151         /**
 152          * Sets the seed used by the internal random number generator.
 153          */
 154         public void setRandomSeed(long seed)
 155         {
 156                 this.random.setSeed(seed);
 157         }
 158
 159         /**
 160          * Retrieves an item with probability proportional to its weight in the
 161          * {@code Map} provided in the input.
 162          */
 163         public T nextItem()
 164         {
 165                 ItemPair item_pair = item_pairs.get(random.nextInt(item_pairs.size()));
 166                 if (random.nextDouble() < item_pair.weight)
 167                         return item_pair.light;
 168                 return item_pair.heavy;
 169         }
 170
 171         /**
 172          * Manages light object/heavy object/light conditional probability tuples.
 173          */
 174         private class ItemPair
 175         {
 176                 T light;
 177                 T heavy;
 178                 double weight;
 179
 180                 private ItemPair(T light, T heavy, double weight)
 181                 {
 182                         this.light = light;
 183                         this.heavy = heavy;
 184                         this.weight = weight;
 185                 }
 186
 187                 @Override
 188         public String toString()
 189                 {
 190                         return String.format("[L:%s, H:%s, %.3f]", light, heavy, weight);
 191                 }
 192         }
 193 }