讀完本文，你不僅學會了演算法套路，還可以順便去 LeetCode 上拿下如下題目：

710.黑名單中的隨機數

-----------

本文講兩道比較有技巧性的資料結構設計題，都是和隨機讀取元素相關的，我們前文 水塘抽樣演算法 也寫過類似的問題。

下面來一道道看。

實現隨機集合

這是力扣第 380 題，看下題目：

就是說就是讓我們實現如下一個類：

class RandomizedSet {public:    /** 如果 val 不存在集合中，則插入並返回 true，否則直接返回 false */     bool insert(int val) {}        /** 如果 val 在集合中，則刪除並返回 true，否則直接返回 false */    bool remove(int val) {}        /** 從集合中等機率地隨機獲得一個元素 */    int getRandom() {}}

本題的難點在於兩點：

2、getRandom 方法返回的元素必須等機率返回隨機元素，也就是說，如果集合裡面有 n 個元素，每個元素被返回的機率必須是 1/n。

HashSet 肯定算一個對吧。雜湊集合的底層原理就是一個大陣列，我們把元素透過雜湊函式對映到一個索引上；如果用拉鍊法解決雜湊衝突，那麼這個索引可能連著一個連結串列或者紅黑樹。

那麼請問對於這樣一個標準的 HashSet，你能否在 O(1) 的時間內實現 getRandom 函式？

其實是不能的，因為根據剛才說到的底層實現，元素是被雜湊函式「分散」到整個數組裡面的，更別說還有拉鍊法等等解決雜湊衝突的機制，基本做不到 O(1) 時間等機率隨機獲取元素。

除了 HashSet，還有一些類似的資料結構，比如雜湊連結串列 LinkedHashSet，我們前文 手把手實現LRU演算法 和 手把手實現LFU演算法 講過這類資料結構的實現原理，本質上就是雜湊表配合雙鏈表，元素儲存在雙鏈表中。

但是，LinkedHashSet 只是給 HashSet 增加了有序性，依然無法按要求實現我們的 getRandom 函式，因為底層用連結串列結構儲存元素的話，是無法在 O(1) 的時間內訪問某一個元素的。

根據上面的分析，對於 getRandom 方法，如果想「等機率」且「在 O(1) 的時間」取出元素，一定要滿足：底層用陣列實現，且陣列必須是緊湊的。

這樣我們就可以直接生成隨機數作為索引，從陣列中取出該隨機索引對應的元素，作為隨機元素。

交換兩個元素必須透過索引進行交換對吧，那麼我們需要一個雜湊表 valToIndex 來記錄每個元素值對應的索引。

有了思路鋪墊，我們直接看程式碼：

class RandomizedSet {public:    // 儲存元素的值    vector<int> nums;    // 記錄每個元素對應在 nums 中的索引    unordered_map<int,int> valToIndex;        bool insert(int val) {        // 若 val 已存在，不用再插入        if (valToIndex.count(val)) {            return false;        }        // 若 val 不存在，插入到 nums 尾部，        // 並記錄 val 對應的索引值        valToIndex[val] = nums.size();        nums.push_back(val);        return true;    }         bool remove(int val) {        // 若 val 不存在，不用再刪除        if (!valToIndex.count(val)) {            return false;        }        // 先拿到 val 的索引        int index = valToIndex[val];        // 將最後一個元素對應的索引修改為 index        valToIndex[nums.back()] = index;        // 交換 val 和最後一個元素        swap(nums[index], nums.back());        // 在陣列中刪除元素 val        nums.pop_back();        // 刪除元素 val 對應的索引        valToIndex.erase(val);        return true;    }        int getRandom() {        // 隨機獲取 nums 中的一個元素        return nums[rand() % nums.size()];    }};

注意 remove(val) 函式，對 nums 進行插入、刪除、交換時，都要記得修改雜湊表 valToIndex，否則會出現錯誤。

至此，這道題就解決了，每個操作的複雜度都是 O(1)，且隨機抽取的元素機率是相等的。

避開黑名單的隨機數

有了上面一道題的鋪墊，我們來看一道更難一些的題目，力扣第 710 題，我來描述一下題目：

給你輸入一個正整數 N，代表左閉右開區間 [0,N)，再給你輸入一個數組 blacklist，其中包含一些「黑名單數字」，且 blacklist 中的數字都是區間 [0,N) 中的數字。

現在要求你設計如下資料結構：

class Solution {public:    // 建構函式，輸入引數    Solution(int N, vector<int>& blacklist) {}        // 在區間 [0,N) 中等機率隨機選取一個元素並返回    // 這個元素不能是 blacklist 中的元素    int pick() {}};

pick 函式會被多次呼叫，每次呼叫都要在區間 [0,N) 中「等機率隨機」返回一個「不在 blacklist 中」的整數。

這應該不難理解吧，比如給你輸入 N = 5, blacklist = [1,3]，那麼多次呼叫 pick 函式，會等機率隨機返回 0, 2, 4 中的某一個數字。

而且題目要求，在 pick 函式中應該儘可能少呼叫隨機數生成函式 rand()。

這句話什麼意思呢，比如說我們可能想出如下拍腦袋的解法：

int pick() {    int res = rand() % N;    while (res exists in blacklist) {        // 重新隨機一個結果        res = rand() % N;    }    return res;}

這個函式會多次呼叫 rand() 函式，執行效率竟然和隨機數相關，不是一個漂亮的解法。

聰明的解法類似上一道題，我們可以將區間 [0,N) 看做一個數組，然後將 blacklist 中的元素移到陣列的最末尾，同時用一個雜湊表進行對映：

根據這個思路，我們可以寫出第一版程式碼（還存在幾處錯誤）：

class Solution {public:    int sz;    unordered_map<int, int> mapping;        Solution(int N, vector<int>& blacklist) {        // 最終陣列中的元素個數        sz = N - blacklist.size();        // 最後一個元素的索引        int last = N - 1;        // 將黑名單中的索引換到最後去        for (int b : blacklist) {            mapping[b] = last;            last--;        }    }};

如上圖，相當於把黑名單中的數字都交換到了區間 [sz, N) 中，同時把 [0, sz) 中的黑名單數字對映到了正常數字。

根據這個邏輯，我們可以寫出 pick 函式：

int pick() {    // 隨機選取一個索引    int index = rand() % sz;    // 這個索引命中了黑名單，    // 需要被對映到其他位置    if (mapping.count(index)) {        return mapping[index];    }    // 若沒命中黑名單，則直接返回    return index;}

這個 pick 函式已經沒有問題了，但是建構函式還有兩個問題。

第一個問題，如下這段程式碼：

int last = N - 1;// 將黑名單中的索引換到最後去for (int b : blacklist) {    mapping[b] = last;    last--;}

我們將黑名單中的 b 對映到 last，但是我們能確定 last 不在 blacklist 中嗎？

比如下圖這種情況，我們的預期應該是 1 對映到 3，但是錯誤地對映到 4：

在對 mapping[b] 賦值時，要保證 last 一定不在 blacklist 中，可以如下操作：

// 建構函式Solution(int N, vector<int>& blacklist) {    sz = N - blacklist.size();    // 先將所有黑名單數字加入 map    for (int b : blacklist) {         // 這裡賦值多少都可以        // 目的僅僅是把鍵存進雜湊表        // 方便快速判斷數字是否在黑名單內        mapping[b] = 666;    }    int last = N - 1;    for (int b : blacklist) {        // 跳過所有黑名單中的數字        while (mapping.count(last)) {            last--;        }        // 將黑名單中的索引對映到合法數字        mapping[b] = last;        last--;    }}

第二個問題，如果 blacklist 中的黑名單數字本身就存在區間 [sz, N) 中，那麼就沒必要在 mapping 中建立對映，比如這種情況：

我們根本不用管 4，只希望把 1 對映到 3，但是按照 blacklist 的順序，會把 4 對映到 3，顯然是錯誤的。

我們可以稍微修改一下，寫出正確的解法程式碼：

class Solution {public:    int sz;    unordered_map<int, int> mapping;        Solution(int N, vector<int>& blacklist) {        sz = N - blacklist.size();        for (int b : blacklist) {            mapping[b] = 666;        }        int last = N - 1;        for (int b : blacklist) {            // 如果 b 已經在區間 [sz, N)            // 可以直接忽略            if (b >= sz) {                continue;            }            while (mapping.count(last)) {                last--;            }            mapping[b] = last;            last--;        }    }    // 見上文程式碼實現    int pick() {}};

至此，這道題也解決了，總結一下本文的核心思想：

1、如果想高效地，等機率地隨機獲取元素，就要使用陣列作為底層容器。

3、對於第二題，陣列中含有「空洞」（黑名單數字），也可以利用雜湊表巧妙處理對映關係，讓陣列在邏輯上是緊湊的，方便隨機取元素。