背景
最近開發個需求,要求是給符合要求的用戶按照情況計算權重,然後隨機抽取幸運用戶獲得獎品。
比如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