权重抽奖实现(数组&TreeMap)

背景

最近开发个需求，要求是给符合要求的用户按照情况计算权重，然后随机抽取幸运用户获得奖品。

比如ABCD用户，权重分别是 1、2、3、4，即概率分别10%、20%、30%、40%

实现一：数组二分查找法

这里的实现原理类似Nacos的负载选择。就是根据权重，得到个数组。

如上述例子，A的区间为[0,1)  B为[1,3)  C为[3,6)  D为[6,10) 

假设获取到的随机数是5.5 那就是落到3和6之间，即选中C。

用数组二分查找就是{0，1，3，6，10} 找到index为2，即选中C。

import java.util.*;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicInteger;

public class WeightRandom<T> {

    // 两个属性 一个是参与活动的集合 泛型 可以是User对象
    private final List<T> items = new ArrayList<>();
    // 权重数组
    private double[] weights;

    public WeightRandom(List<ItemWithWeight<T>> itemsWithWeight) {
        this.calWeights(itemsWithWeight);
    }

    /**
     * 计算权重，初始化或者重新定义权重时使用
     *
     */
    public void calWeights(List<ItemWithWeight<T>> itemsWithWeight) {
        items.clear();

        // 计算权重总和
        double originWeightSum = 0;
        // 遍历所有参与的对象 
        for (ItemWithWeight<T> itemWithWeight : itemsWithWeight) {
            double weight = itemWithWeight.getWeight();
            if (weight <= 0) {
                continue;
            }

            items.add(itemWithWeight.getItem());
            // 是否无穷大  无穷大取个10000
            if (Double.isInfinite(weight)) {
                weight = 10000.0D;
            }
            // 是否无权重  这里设置取1
            if (Double.isNaN(weight)) {
                weight = 1.0D;
            }
            // 将所有的权重值相加
            originWeightSum += weight;
        }

        // 计算每个item的实际权重比例
        double[] actualWeightRatios = new double[items.size()];
        int index = 0;
        for (ItemWithWeight<T> itemWithWeight : itemsWithWeight) {
            double weight = itemWithWeight.getWeight();
            if (weight <= 0) {
                continue;
            }
            // 得到权重占比 放置数组中
            actualWeightRatios[index++] = weight / originWeightSum;
        }

        // 计算每个item的权重范围
        // 权重范围起始位置
        weights = new double[items.size()];
        double weightRangeStartPos = 0;
        for (int i = 0; i < index; i++) {
            weights[i] = weightRangeStartPos + actualWeightRatios[i];
            weightRangeStartPos += actualWeightRatios[i];
        }
        // 得到一个数组 类如{0,3,5,9,11,16....}
    }

    /**
     * 基于权重随机算法选择
     *
     */
    public T choose() {
        // 得到一个随机数
        double random = ThreadLocalRandom.current().nextDouble();
        // 用二分查找 找到index
        int index = Arrays.binarySearch(weights, random);
        if (index < 0) {
            index = -index - 1;
        } else {
            return items.get(index);
        }

        if (index < weights.length && random < weights[index]) {
            return items.get(index);
        }

        // 通常不会走到这里，为了保证能得到正确的返回，这里随便返回一个
        return items.get(0);
    }

    public static class ItemWithWeight<T> {
        T item;
        double weight;

        public ItemWithWeight() {
        }

        public ItemWithWeight(T item, double weight) {
            this.item = item;
            this.weight = weight;
        }

        public T getItem() {
            return item;
        }

        public void setItem(T item) {
            this.item = item;
        }

        public double getWeight() {
            return weight;
        }

        public void setWeight(double weight) {
            this.weight = weight;
        }
    }

    public static void main(String[] args) {
        // for test
        int sampleCount = 1_000_000;

        ItemWithWeight<String> server1 = new ItemWithWeight<>("server1", 1.0);
        ItemWithWeight<String> server2 = new ItemWithWeight<>("server2", 3.0);
        ItemWithWeight<String> server3 = new ItemWithWeight<>("server3", 2.0);

        WeightRandom<String> weightRandom = new WeightRandom<>(Arrays.asList(server1, server2, server3));

        // 统计 (这里用 AtomicInteger 仅仅是因为写起来比较方便，这是一个单线程测试)
        Map<String, AtomicInteger> statistics = new HashMap<>();

        // 模拟一次
        System.out.println(weightRandom.choose());

        for (int i = 0; i < sampleCount; i++) {
            statistics
                    .computeIfAbsent(weightRandom.choose(), (k) -> new AtomicInteger())
                    .incrementAndGet();
        }

        // 模拟10w次的概率分布
        statistics.forEach((k, v) -> {
            double hit = (double) v.get() / sampleCount;
            System.out.println(k + ", hit:" + hit);
        });
    }
}

 

实现二：TreeMap二叉树

这里运用了TreeMap底层有序的二叉分布，每个节点的左边都是小于当前节点，右边都是大于当前节点的。

按照树的层级遍历，效率与数组二分查找差不多。

import cn.hutool.core.lang.Pair;
import com.google.common.base.Preconditions;
import lombok.extern.slf4j.Slf4j;


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

@Slf4j
public class WeightRandomByTreeMap<K,V extends Number> {
	private final TreeMap<Double, K> weightMap = new TreeMap<>();


	public WeightRandomByTreeMap(List<Pair<K, V>> list) {
		Preconditions.checkNotNull(list, "list can NOT be null!");
		for (Pair<K, V> pair : list) {
			Preconditions.checkArgument(pair.getValue().doubleValue() > 0, String.format("非法权重值：pair=%s", pair));

			//统一转为double
			double lastWeight = this.weightMap.size() == 0 ? 0 : this.weightMap.lastKey();
			//权重累加
			this.weightMap.put(pair.getValue().doubleValue() + lastWeight, pair.getKey());
		}
	}

	public K random() {
		double randomWeight = this.weightMap.lastKey() * Math.random();
        // 核心逻辑 tailMap得到大于等于随机数的有序集合
		SortedMap<Double, K> tailMap = this.weightMap.tailMap(randomWeight, false);
        // 取第一个 就是抽中的用户了
		return this.weightMap.get(tailMap.firstKey());
	}

}

两种实现逻辑都是先按照权重给用户划分区域，然后生成随机数看落在哪个区域即选中。

 

 

 

 

参考：

https://www.cnblogs.com/waterystone/p/5708063.html