給我 O(1) 的時間，我可以刪除/查詢陣列中的任意元素
給我 O(1) 的時間，我可以刪除/查詢陣列中的任意元素
讀完本文，你不僅學會了演算法套路，還可以順便去 LeetCode 上拿下如下題目：
-----------
本文講兩道比較有技巧性的資料結構設計題，都是和隨機讀取元素相關的，我們前文 水塘抽樣演算法 也寫過類似的問題。
這寫問題的一個技巧點在於，如何結合雜湊表和陣列，使得陣列的刪除操作時間複雜度也變成 O(1)？
下面來一道道看。
實現隨機集合
這是力扣第 380 題，看下題目：
就是說就是讓我們實現如下一個類：
class RandomizedSet { public: /** 如果 val 不存在集合中，則插入並返回 true，否則直接返回 false */ bool insert(int val) {} /** 如果 val 在集合中，則刪除並返回 true，否則直接返回 false */ bool remove(int val) {} /** 從集合中等概率地隨機獲得一個元素 */ int getRandom() {} }
本題的難點在於兩點：
1、插入，刪除，獲取隨機元素這三個操作的時間複雜度必須都是 O(1)。
2、getRandom 方法返回的元素必須等概率返回隨機元素，也就是說，如果集合裡面有 n 個元素，每個元素被返回的概率必須是 1/n。
我們先來分析一下：對於插入，刪除，查詢這幾個操作，哪種資料結構的時間複雜度是 O(1)？
HashSet 肯定算一個對吧。雜湊集合的底層原理就是一個大陣列，我們把元素通過雜湊函式對映到一個索引上；如果用拉鍊法解決雜湊衝突，那麼這個索引可能連著一個連結串列或者紅黑樹。
那麼請問對於這樣一個標準的 HashSet，你能否在 O(1) 的時間內實現 getRandom 函式？
其實是不能的，因為根據剛才說到的底層實現，元素是被雜湊函式「分散」到整個數組裡面的，更別說還有拉鍊法等等解決雜湊衝突的機制，基本做不到 O(1) 時間等概率隨機獲取元素。
除了 HashSet，還有一些類似的資料結構，比如雜湊連結串列 LinkedHashSet，我們前文 手把手實現LRU演算法 和 手把手實現LFU演算法 講過這類資料結構的實現原理，本質上就是雜湊表配合雙鏈表，元素儲存在雙鏈表中。
但是，LinkedHashSet 只是給 HashSet 增加了有序性，依然無法按要求實現我們的 getRandom 函式，因為底層用連結串列結構儲存元素的話，是無法在 O(1) 的時間內訪問某一個元素的。
根據上面的分析，對於 getRandom 方法，如果想「等概率」且「在 O(1) 的時間」取出元素，一定要滿足：底層用陣列實現，且陣列必須是緊湊的。
這樣我們就可以直接生成隨機數作為索引，從陣列中取出該隨機索引對應的元素，作為隨機元素。
但如果用陣列儲存元素的話，插入，刪除的時間複雜度怎麼可能是 O(1) 呢？
可以做到！對陣列尾部進行插入和刪除操作不會涉及資料搬移，時間複雜度是 O(1)。
所以，如果我們想在 O(1) 的時間刪除陣列中的某一個元素 val，可以先把這個元素交換到陣列的尾部，然後再 pop 掉。
交換兩個元素必須通過索引進行交換對吧，那麼我們需要一個雜湊表 valToIndex 來記錄每個元素值對應的索引。
PS：我認真寫了 100 多篇原創，手把手刷 200 道力扣題目，全部發布在 labuladong的演算法小抄，持續更新。建議收藏，按照我的文章順序刷題，掌握各種演算法套路後投再入題海就如魚得水了。
有了思路鋪墊，我們直接看程式碼：
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 中的某一個數字。
PS：我認真寫了 100 多篇原創，手把手刷 200 道力扣題目，全部發布在 labuladong的演算法小抄，持續更新。建議收藏，按照我的文章順序刷題，掌握各種演算法套路後投再入題海就如魚得水了。
而且題目要求，在 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、如果想高效地，等概率地隨機獲取元素，就要使用陣列作為底層容器。
2、如果要保持陣列元素的緊湊性，可以把待刪除元素換到最後，然後 pop 掉末尾的元素，這樣時間複雜度就是 O(1) 了。當然，我們需要額外的雜湊表記錄值到索引的對映。
3、對於第二題，陣列中含有「空洞」（黑名單數字），也可以利用雜湊表巧妙處理對映關係，讓陣列在邏輯上是緊湊的，方便隨機取元素。
相關推薦：
＿＿＿＿＿＿＿＿＿＿＿＿＿
點選 我的主頁 看更多優質文章。