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實現
權重隨機生成