java实现权重随机算法

package com.jmmq.load.jim.algorithm;


import java.math.BigDecimal;
import java.util.*;

/**
 * 权重随机算法
 *  主要思路就是数据按照权重进行位置区分，
 *  然后使用随机数判断落入的区间，这样就进行了随机
 * ---------------------------
 *      一等奖  0.05
 *      二等奖  0.35
 *      三等奖  0.6
 * ---------------------------
 *  ↓ 转换为  权重边界 上包含 0,0.05 -> 一等奖  0.40 -> 二等奖  1.0 三等奖
 *  --------------------------
 *     一等奖  0 ~ 0.05
 *     二等奖  0.05 ~ 0.40
 *     三等奖  0.40 ~ 1.0
 * ---------------------------
 *    这里要求各个选项的概率和为1.0
 *
 *    参考资料： https://www.zifangsky.cn/1545.html
 *             https://blog.csdn.net/zj20142213/article/details/90344388
 */
public class WeightRandomPrc {

    public static void main(String[] args) {
        WeightRandomItem[] items = {
                new WeightRandomItem("一等奖","0.05"),
                new WeightRandomItem("二等奖","0.15"),
                new WeightRandomItem("三等奖","0.3"),
                new WeightRandomItem("纪念奖","0.5")

        };

        // 测试1
        for(int i=0; i < 10; i++) {
            System.out.println("第"+ (i+1) + "次：" + weightRandom(items));
        }

        System.out.println("---------------------------------------------------");

        // 优化1测试
        WeightRandom randomHandler =  WeightRandom.getInstance(items);
        for(int i=0; i < 10; i++) {
            System.out.println("第"+ (i+1) + "次：" + randomHandler.radomNext());
        }

        System.out.println("---------------------------------------------------");

        // TreeMap
        WeightRandomWithTreeMap weightRandomWithTreeMap = new WeightRandomWithTreeMap(items);
        for(int i=0; i < 10; i++) {
            System.out.println("第"+ (i+1) + "次：" + weightRandomWithTreeMap.randomNext());
        }
    }

    static class WeightRandomItem{
        // 名称
        private String name;
        // 权重比例
        private String probability;

        public WeightRandomItem(String name, String probability) {
            this.name = name;
            this.probability = probability;
        }

        public String getName() {
            return name;
        }

        public void setName(String name) {
            this.name = name;
        }

        public String getProbability() {
            return probability;
        }

        public void setProbability(String probability) {
            this.probability = probability;
        }
    }

    public static String weightRandom(WeightRandomItem[] items){

        // 校验值的和是否为1.0
        List<WeightRandomItem> itemList = Arrays.asList(items);
        BigDecimal sum = itemList.stream()
                .map(i -> {return new BigDecimal(i.probability);})
                .reduce(BigDecimal.ZERO, BigDecimal::add);

        if(new BigDecimal("1").compareTo(sum) != 0) {
            System.out.println("概率之和必须等于1");
            return "";
        }

        // 计算每一个值的权重边界
        /**
         * 这里拿到的结果  将数据转换为下标
         *    0.05      0.40    1.0
         *   一等奖     二等奖    三等奖
         */
        BigDecimal[] weight = new BigDecimal[itemList.size()];
        BigDecimal sumP = BigDecimal.ZERO;
        for(int i=0; i< itemList.size(); i++){
            sumP = sumP.add(new BigDecimal(itemList.get(i).getProbability()));
            weight[i] = sumP;
        }

        // 随机数
        BigDecimal point = new BigDecimal(new Random().nextDouble());
//        System.out.println("随机值：" + point.toString());
        // 二分查找point所属范围
        int low = 0;
        int high = weight.length -1;
        int middle = (high + low) / 2;
        while(
                // middle != low && middle != high
                high - low > 1
        ){
            BigDecimal middleValue = weight[middle];
            if(point.compareTo(middleValue) == 0) {
                return items[middle].getName();
            } else if (point.compareTo(middleValue) > 0){
                low = middle;
                middle = (high + low) / 2;
            } else {
                high = middle;
                middle = (high + low) / 2;
            }
//            System.out.println("low:" + low + ",high:" + high + ",middle:" + middle);
        }

        BigDecimal highValue = weight[high];
        BigDecimal lowValue = weight[low];

        if(highValue.compareTo(point) < 0  ){
            return items[low].getName();
        }

        return items[high].getName();
    }

    // 优化,将上面的方法抽象一下
    static class WeightRandom{
        // 权重数组
        private  WeightRandomItem[] items;
        // 权重边界数组
        private BigDecimal[] weight;

        private static volatile WeightRandom weightRandom = null;

        public static WeightRandom getInstance(WeightRandomItem[] items){
            if(weightRandom == null){
                synchronized (WeightRandom.class){
                    if(weightRandom==null){
                        weightRandom = new WeightRandom(items);
                    }
                }
            }
            return weightRandom;
        }

        private WeightRandom(WeightRandomItem[] items) {
            this.items = items;
            this.init();
        }

        private String radomNext(){
            // 随机数
            BigDecimal point = new BigDecimal(new Random().nextDouble());

            int low = 0;
            int high = weight.length -1;
            int middle = (high + low) / 2;
            while(high - low > 1){
                BigDecimal middleValue = weight[middle];
                if(point.compareTo(middleValue) == 0) {
                    return items[middle].getName();
                } else if (point.compareTo(middleValue) > 0){
                    low = middle;
                    middle = (high + low) / 2;
                } else {
                    high = middle;
                    middle = (high + low) / 2;
                }
            }

            BigDecimal highValue = weight[high];
            BigDecimal lowValue = weight[low];

            if(highValue.compareTo(point) < 0  ){
                return items[low].getName();
            }

            return items[high].getName();
        }


        private void clear(){
            this.setItems(null);
            this.setWeight(null);
        }

        // 组装权重边界数组、校验权重比例和是否为1
        private void init(){
            BigDecimal sum = BigDecimal.ZERO;
            weight = new BigDecimal[items.length];
            for(int i=0; i<items.length; i++){
                sum = sum.add(new BigDecimal(items[i].getProbability()));
                weight[i] = sum;
            }

            if(sum.compareTo(new BigDecimal("1")) !=0){
                throw new RuntimeException("probability sum must be 1.0" );
            }
            System.out.println("init sucess");
        }

        private WeightRandomItem[] getItems() {
            return items;
        }

        private void setItems(WeightRandomItem[] items) {
            this.items = items;
        }

        private BigDecimal[] getWeight() {
            return weight;
        }

        private void setWeight(BigDecimal[] weight) {
            this.weight = weight;
        }
    }

    // 利用java的TreeMap 、这个当然还可以优化代码扩展适用性
    static class WeightRandomWithTreeMap{

        // 权重比例、值
        private TreeMap<BigDecimal, String> weightMap = new TreeMap<BigDecimal, String>();

        private WeightRandomItem[] items;

        public WeightRandomWithTreeMap(WeightRandomItem[] items) {
            this.items = items;
            this.init();
        }

        private void init(){
            BigDecimal sum = BigDecimal.ZERO;
            for(int i=0; i<items.length; i++){
                sum = sum.add(new BigDecimal(items[i].getProbability()));
                weightMap.put(sum, items[i].getName());
            }

            if(sum.compareTo(new BigDecimal("1")) !=0){
                throw new RuntimeException("probability sum must be 1.0" );
            }
        }

        private String randomNext(){
            BigDecimal point = new BigDecimal(new Random().nextDouble());
            SortedMap<BigDecimal, String> tailMap = this.weightMap.tailMap(point, true);
            return this.weightMap.get(tailMap.firstKey());
        }

    }
}

第1次：二等奖
第2次：三等奖
第3次：二等奖
第4次：纪念奖
第5次：三等奖
第6次：三等奖
第7次：二等奖
第8次：纪念奖
第9次：纪念奖
第10次：三等奖
第1次：纪念奖
第2次：二等奖
第3次：纪念奖
第4次：三等奖
第5次：纪念奖
第6次：纪念奖
第7次：纪念奖
第8次：三等奖
第9次：三等奖
第10次：纪念奖
第1次：三等奖
第2次：三等奖
第3次：纪念奖
第4次：二等奖
第5次：纪念奖
第6次：纪念奖
第7次：三等奖
第8次：一等奖
第9次：三等奖
第10次：纪念奖