-----------我最近在 LeetCode 上做到兩道非常有意思的題目,382 和 398 題,關於水塘抽樣算法(Reservoir Sampling),本質上是一種隨機概率算法,解法應該說會者不難,難者不會。
我第一次見到這個算法問題是谷歌的一道算法題:給你一個未知長度的鏈表,請你設計一個算法,只能遍歷一次,隨機地返回鏈表中的一個節點。
這裏說的隨機是均勻隨機(uniform random),也就是說,如果有 n
個元素,每個元素被選中的概率都是 1/n
,不可以有統計意義上的偏差。
一般的想法就是,我先遍歷一遍鏈表,得到鏈表的總長度 n
,再生成一個 [1,n]
之間的隨機數爲索引,然後找到索引對應的節點,不就是一個隨機的節點了嗎?
但題目說了,只能遍歷一次,意味着這種思路不可行。題目還可以再泛化,給一個未知長度的序列,如何在其中隨機地選擇 k
個元素?想要解決這個問題,就需要著名的水塘抽樣算法了。
PS:我認真寫了 100 多篇原創,手把手刷 200 道力扣題目,全部發布在labuladong的算法小抄,持續更新。建議收藏,按照我的文章順序刷題,掌握各種算法套路後投再入題海就如魚得水了。
算法實現
先解決只抽取一個元素的問題,這個問題的難點在於,隨機選擇是「動態」的,比如說你現在你有 5 個元素,你已經隨機選取了其中的某個元素 a
作爲結果,但是現在再給你一個新元素 b
,你應該留着 a
還是將 b
作爲結果呢,以什麼邏輯選擇 a
和 b
呢,怎麼證明你的選擇方法在概率上是公平的呢?
先說結論,當你遇到第 i
個元素時,應該有 1/i
的概率選擇該元素,1 - 1/i
的概率保持原有的選擇。看代碼容易理解這個思路:
/* 返回鏈表中一個隨機節點的值 */
int getRandom(ListNode head) {
Random r = new Random();
int i = 0, res = 0;
ListNode p = head;
// while 循環遍歷鏈表
while (p != null) {
// 生成一個 [0, i) 之間的整數
// 這個整數等於 0 的概率就是 1/i
if (r.nextInt(++i) == 0) {
res = p.val;
}
p = p.next;
}
return res;
}
對於概率算法,代碼往往都是很淺顯的,但是這種問題的關鍵在於證明,你的算法爲什麼是對的?爲什麼每次以 1/i
的概率更新結果就可以保證結果是平均隨機(uniform random)?
證明:假設總共有 n
個元素,我們要的隨機性無非就是每個元素被選擇的概率都是 1/n
對吧,那麼對於第 i
個元素,它被選擇的概率就是:
第 i
個元素被選擇的概率是 1/i
,第 i+1
次不被替換的概率是 1 - 1/(i+1)
,以此類推,相乘就是第 i
個元素最終被選中的概率,就是 1/n
。
因此,該算法的邏輯是正確的。
同理,如果要隨機選擇 k
個數,只要在第 i
個元素處以 k/i
的概率選擇該元素,以 1 - k/i
的概率保持原有選擇即可。代碼如下:
/* 返回鏈表中 k 個隨機節點的值 */
int[] getRandom(ListNode head, int k) {
Random r = new Random();
int[] res = new int[k];
ListNode p = head;
// 前 k 個元素先默認選上
for (int j = 0; j < k && p != null; j++) {
res[j] = p.val;
p = p.next;
}
int i = k;
// while 循環遍歷鏈表
while (p != null) {
// 生成一個 [0, i) 之間的整數
int j = r.nextInt(++i);
// 這個整數小於 k 的概率就是 k/i
if (j < k) {
res[j] = p.val;
}
p = p.next;
}
return res;
}
對於數學證明,和上面區別不大:
因爲雖然每次更新選擇的概率增大了 k
倍,但是選到具體第 i
個元素的概率還是要乘 1/k
,也就回到了上一個推導。
PS:我認真寫了 100 多篇原創,手把手刷 200 道力扣題目,全部發布在labuladong的算法小抄,持續更新。建議收藏,按照我的文章順序刷題,掌握各種算法套路後投再入題海就如魚得水了。
拓展延伸
以上的抽樣算法時間複雜度是 O(n),但不是最優的方法,更優化的算法基於幾何分佈(geometric distribution),時間複雜度爲 O(k + klog(n/k))。由於涉及的數學知識比較多,這裏就不列出了,有興趣的讀者可以自行搜索一下。
還有一種思路是基於「Fisher–Yates 洗牌算法」的。隨機抽取 k
個元素,等價於對所有元素洗牌,然後選取前 k
個。只不過,洗牌算法需要對元素的隨機訪問,所以只能對數組這類支持隨機存儲的數據結構有效。
另外有一種思路也比較有啓發意義:給每一個元素關聯一個隨機數,然後把每個元素插入一個容量爲 k
的二叉堆(優先級隊列)按照配對的隨機數進行排序,最後剩下的 k
個元素也是隨機的。
這個方案看起來似乎有點多此一舉,因爲插入二叉堆需要 O(logk) 的時間複雜度,所以整個抽樣算法就需要 O(nlogk) 的複雜度,還不如我們最開始的算法。但是,這種思路可以指導我們解決加權隨機抽樣算法,權重越高,被隨機選中的概率相應增大,這種情況在現實生活中是很常見的,比如你不往遊戲裏充錢,就永遠抽不到皮膚。
最後,我想說隨機算法雖然不多,但其實很有技巧的,讀者不妨思考兩個常見且看起來很簡單的問題:
1、如何對帶有權重的樣本進行加權隨機抽取?比如給你一個數組 w
,每個元素 w[i]
代表權重,請你寫一個算法,按照權重隨機抽取索引。比如 w = [1,99]
,算法抽到索引 0 的概率是 1%,抽到索引 1 的概率是 99%。
2、實現一個生成器類,構造函數傳入一個很長的數組,請你實現 randomGet
方法,每次調用隨機返回數組中的一個元素,多次調用不能重複返回相同索引的元素。要求不能對該數組進行任何形式的修改,且操作的時間複雜度是 O(1)。