528. Random Pick with Weight

題目描述

給定一個正整數陣列 w，其中 w[i] 代表索引 i 的權重（weight）。實作 pickIndex() 函式，依照權重按比例隨機回傳一個索引。

即索引 i 被選中的概率為 w[i] / sum(w)。

範例：w = [1, 3]

• 索引 0 的概率為 1/4 = 25%
• 索引 1 的概率為 3/4 = 75%

解題思路

核心思路：前綴和 + 二分搜尋

概率轉化：將每個索引的權重視為「號碼牌的數量」。若 w = [1, 3]，可以想像共有 4 張號碼牌（1 張給索引 0，3 張給索引 1），隨機抽取一張。

前綴和建構累積概率：

• 建立前綴和陣列 prefix，其中 prefix[i] = w[0] + w[1] + ... + w[i]
• prefix 的最後一個值為所有權重的總和 total

隨機選取：

1. 隨機生成一個整數 target，範圍為 [1, total]。
2. 在前綴和陣列中找到第一個大於等於 target 的位置（即 lower_bound），該位置就是被選中的索引。

直觀解釋：前綴和陣列將 [1, total] 的整數區間劃分為若干段，索引 i 對應的段為 (prefix[i-1], prefix[i]]，段的長度恰好為 w[i]。隨機數落在哪一段就選哪個索引，選中概率與段長成正比，即與 w[i] 成正比。

為什麼是 [1, total] 而非 [0, total-1]：使用 1 到 total 的閉區間配合 lower_bound（找第一個 ≥ target 的位置），可以統一處理所有情況，避免 off-by-one 錯誤。

時間複雜度

• 建構子：O(n)，建立前綴和陣列需線性遍歷。
• pickIndex：O(log n)，每次生成隨機數後用二分搜尋定位索引。

空間複雜度

O(n)，儲存前綴和陣列。若 n 很大但 pickIndex 呼叫次數也很多，O(log n) per pick 優於 O(n) per pick（線性掃描）。

替代方案

方案一：線性掃描

每次 pickIndex 時，生成 [1, total] 的隨機數，然後線性掃描前綴和陣列，找到第一個大於等於隨機數的位置。

• 優點：不需要二分搜尋，程式碼更簡單，建構仍 O(n)。
• 缺點：每次 pickIndex 時間 O(n)，若呼叫次數 q 很大，總時間 O(n + q×n) 遠不如 O(n + q×log n)。

方案二：儲存展開的陣列（Reservoir Sampling 思路）

預先建立一個包含 total 個元素的陣列，每個索引 i 重複出現 w[i] 次；pickIndex 直接隨機選取陣列中一個位置。

• 優點：pickIndex 時間 O(1)，極為簡單。
• 缺點：空間 O(total)，若權重值很大（如 10^5 個索引每個權重 10^5），空間無法接受；建構時間也是 O(total)。

方案三：Alias Method（別名採樣）

一種 O(n) 預處理後 O(1) 採樣的演算法，適合固定分佈的高頻採樣。

• 優點：採樣 O(1)，建構 O(n)，是理論上最優的隨機採樣方案。
• 缺點：實作複雜（需要 Vose's Alias Method），面試中極少要求；前綴和 + 二分搜尋已足夠應對大多數場景。

延伸問題

1. 動態權重更新：若 w[i] 可以被修改，前綴和方案每次更新需要 O(n) 重建，效率低。可以改用線段樹（Segment Tree）維護區間和，支援 O(log n) 的點更新和 O(log n) 的隨機採樣（在線段樹上做隨機遊走）。

2. 帶拒絕採樣的均勻分佈：若權重不是整數而是浮點數（如概率），如何處理？可以將浮點概率乘以一個大整數轉為近似整數，或改用連續均勻分佈（uniform_real_distribution）配合浮點前綴和做二分搜尋。

3. 在線流式採樣（Reservoir Sampling）：若輸入是串流（不知道總數），如何保持均勻採樣？Reservoir Sampling（蓄水池採樣）算法可以在 O(n) 時間 O(1) 空間內從 n 個元素中均勻隨機選取一個，且不需要事先知道 n 的大小。