給一串正整數權重 w ,w[i] 代表 i 的權重。然後要寫一個 pickIndex 方法,以加權的機率抽出一個 i 。
雖然說 Python 的 random 模組有方法可以模擬做加權抽樣,但是速度不夠快,以致於超時,所以要減少 random 的使用,靠自己寫更有效率的代碼。
(這也是很少見的專門模組效率會遜於自撰的)
加權抽樣若以 1 為單位,把所有 index 放入一個假想的箱子,權重為 1 放1個,權重為 3 放3個,最後從中隨機抽取一個出來。
抽中 i 的機率為 w[i] / total(總權重) ,所有的機率加起來會等於 1 。若把機率從頭逐一累加,保存累進序列,這序列以權重把[0,1]之間切割區間。
用 random.random() ,可以得到一個 [0, 1] 之間的 float ,然後我們要用這個 [0, 1]之間的 random數 去對應出一個 i。
以輸入[1,4]為例,對照累進機率,若 random數 位於 [0, 0.2],那代表抽到 0 。若 random數 位於 [0.2, 1],那代表抽到 1。
要找出 random數 落於哪個區間,可以用二分法去搜尋,預期 O(logN) Time 。(不知道random的方法是用到 O(N) Time嗎?至少二分法比O(N)快)
如果剛好落在區間界線(如 == 0.2)要算前面的 i 還是後面的 i ?其實都沒差,根據「積分」,剛好落在 0.2 的機率為 0 ,所以不會影響到權重。
Python
class Solution:
def __init__(self, w: List[int]):
total = sum(w)
self.w = []
accu = 0
for i in range(1, len(w)):
accu += w[i] / total
self.w.append(accu)
def pickIndex(self) -> int:
l, r = 0, len(self.w)
t = random.random()
while l < r:
m = (r - l)//2 + l
if self.w[m] < t:
l = m + 1
else:
r = m
return l
C#
public class Solution {
double[] ws;
Random random = new Random();
public Solution(int[] w) {
var total = w.Sum() * 1.0;
ws = new double[w.Length];
ws[0] = w[0] / total;
for(int i = 1; i < w.Length; i++)
{
ws[i] = ws[i-1] + w[i] / total;
}
}
public int PickIndex() {
var t = random.NextDouble();
int l = 0; int r = ws.Length;
while(l < r)
{
int m = (r-l)/2 + l;
if(ws[m] < t)
l = m + 1;
else
r = m;
}
return l;
}
}