java实现随机权重算法

前言

适用场景：一个箱子中有3个元素A，B，C，抽到A的概率为50%，B概率为20%，C概率为30%。我们可以给A，B，C各附加一个权重值，如50,20,30。

实现原理

以上面的A,B,C为例，A的权重区间为[0,50)，B的区间为[50,70), C区间为[70,100)，生成一个100之内的随机值，看落在哪个区间内，就是那个元素。

代码实现

import java.util.SortedMap;
import java.util.TreeMap;


public class WeightRandom<T> {

  private final TreeMap<Integer, T> weightMap;

  /**
   * 创建权重随机获取器
   *
   * @param <T> 权重随机获取的对象类型
   * @return {@link WeightRandom}
   */
  public static <T> WeightRandom<T> create() {
    return new WeightRandom<>();
  }

  /**
   * 构造
   */
  private WeightRandom() {
    weightMap = new TreeMap<>();
  }

  /**
   * 增加对象
   *
   * @param obj 对象
   * @param weight 权重
   */
  public void add(T obj, int weight) {
    if (weight > 0) {
      int lastWeight = (this.weightMap.size() == 0) ? 0 : this.weightMap.lastKey();
      this.weightMap.put(weight + lastWeight, obj);// 权重累加
    }
  }

  /**
   * 清空权重表
   */
  public void clear() {
    this.weightMap.clear();
  }

  /**
   * 下一个随机对象
   *
   * @return 随机对象
   */
  public T next() {
    int randomWeight = (int) (this.weightMap.lastKey() * Math.random());
    SortedMap<Integer, T> tailMap = this.weightMap.tailMap(randomWeight, false);
    return this.weightMap.get(tailMap.firstKey());
  }

}

使用

import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;

public class Client {

  public static void main(String[] args) {
    WeightRandom<String> weightRandom = WeightRandom.create();
    weightRandom.add("A", 50);
    weightRandom.add("B", 20);
    weightRandom.add("C", 30);
    testRate(weightRandom);
  }

  private static void testRate(WeightRandom<String> weightRandom) {
    Map<String, Integer> countMap = new HashMap<>();
    int total = 100_000_000;
    for (int i = 0; i < total; i++) {
      String randomKey = weightRandom.next();
      countMap.put(randomKey, countMap.getOrDefault(randomKey, 0) + 1);
    }
    for (Entry<String, Integer> entry : countMap.entrySet()) {
      System.out.println(String.format("%s count: %s,freq: %s", entry.getKey(), entry.getValue(),
          1.0 * entry.getValue() / total));
    }
  }
}

输出结果为

A count: 49994992,freq: 0.49994992
B count: 20005474,freq: 0.20005474
C count: 29999534,freq: 0.29999534

频率和概率基本相同。

参考

权重随机算法的java实现
权重随机生成