np.random.choice 具有很大的概率數組

Question

我知道我們可以使用概率數組來選擇 function，但我的問題是它如何適用於大 arrays。假設我想要 0-65535 之間的 1000 個隨機數。 我們如何定義概率數組，使其對於小於 1000 的數字具有 p=0.4，對於 rest 具有 p=0.6？

我試圖將數字范圍傳遞給選項 function，但顯然，它不是那樣工作的。

Answer 1

從文檔中，參數p的每個元素都給出了a中對應元素的概率。

由於p和a需要具有相同的大小，因此創建一個與 a 大小相同a p ：

a = np.arange(65536)
n_elem = len(a)

p = np.zeros_like(a, dtype=float)

現在，找到a小於1000的所有元素，並將這些索引的p設置為 0.4 除以小於 1000 的元素數。對於這種情況，您可以對該計算進行硬編碼，因為您知道arange的哪些元素更少大於 1000：

p[:1000] = 0.4 / 1000
p[1000:] = 0.6 / 64536

對於a不是從arange派生的一般情況，您可以這樣做：

lt1k = a < 1000
n_lt1k = lt1k.sum()

p[lt1k] = 0.4 / n_lt1k
p[~lt1k] = 0.6 / (n_elem - n_lt1k)

請注意， p的總和必須為1 ：

assert np.allclose(p.sum(), 1.0)

現在在choice中使用a和p ：

selection = np.random.choice(a, size=(1000,), p=p)

為了驗證選擇值 < 1000 的概率是 40%，我們可以檢查有多少小於 1000：

print((selection < 1000).sum() / len(selection)) # should print a number close to 0.4

Answer 2

另一種方法是將其視為兩種分布的混合：一種以概率 = 0.4 均勻地從 {0..999} 抽取，另一種以概率 = 0.6 均勻地從 {1000..65535} 抽取。

對混合組件使用choice是有道理的，但隨后我會使用其他東西來繪制值，因為當概率傳遞給choice時，它會在每次調用時 O( len(p) ) 工作以轉換它們。 Generator.integers應該更有效，因為它可以直接對統一值進行采樣。

把這些放在一起，我建議使用類似的東西：

import numpy as np

rng = np.random.default_rng()

n = 1000
splits = np.array([0, 1000, 65536])

# draw weighted mixture components
s = rng.choice(2, n, p=[0.4, 0.6])
# draw uniform values according to component
result = rng.integers(splits[s], splits[s+1])

您可以通過評估np.mean(result < 1000)並檢查它是否“接近”0.4 來驗證這是從正確的分布中提取的。 其方差約為0.4*0.6 / n ，因此，對於n=1000 ， [0.37, 0.43] 中的值應該在 95% 的時間內可見。

當max(splits) - min(splits)變大時，此方法應該保持快速，而 Pranav 直接使用choice的解決方案會變慢。