带权重值的随机选择算法实现

带权重值的随机选择算法实现

  •  10个月前
  •  843
  •  Java 
import java.util.HashMap;
import java.util.Map;
import java.util.SortedMap;
import java.util.TreeMap;
import java.util.stream.IntStream;

public class WeightRandom<K, V extends Number> {

    private TreeMap<Double, K> weightMap = new TreeMap<>();

    public WeightRandom(Map<K, V> map) {
        for (K key : map.keySet()) {
            double weight = map.get(key).doubleValue();
            if (weight <= 0) continue;
            double lastWeight = this.weightMap.size() == 0 ? 0 : this.weightMap.lastKey();
            this.weightMap.put(weight + lastWeight, key);
        }
    }

    public K random() {
        double randomWeight = this.weightMap.lastKey() * Math.random();
        SortedMap<Double, K> tailMap = this.weightMap.tailMap(randomWeight, false);
        return this.weightMap.get(tailMap.firstKey());
    }


    public static void main(String[] args) {
        HashMap<String, Integer> map = new HashMap<>();
        map.put("A", 1);
        map.put("B", 2);
        map.put("C", 4);

        WeightRandom<String, Integer> weightRandom = new WeightRandom<>(map);
        HashMap<String, Integer> result = new HashMap<>();
        // 执行100000次
        IntStream.range(0, 100000).forEach(i -> {
            String random = weightRandom.random();
            result.merge(random, 1, (x, y) -> x + y);
        });
        System.out.println(result);
        // 输出结果   {A=14390, B=28730, C=56880}
    }

}


扫一扫分享到微信

已有 条评论
写评论