本题有两个难点:随机返回元素和用$O(1)$时间删除数组中元素,解决思路如下:
- 随机返回元素,使每个元素有相同的概率返回
如果只用哈希表则不能等概率返回其中的每个数值,如果数值是保存在数组中的话,那么就可以通过
rand()
函数轻易做到等概率返回。我们开一个数组来存数值,并在哈希表中储存每个数值及其在数组中的下标。 - 在$O(1)$时间复杂度内删除数组中元素
如果直接
remove
数组中的元素,那么数组中后面的元素会向前移动,时间复杂度为$O(n)$。我们需要先从哈希表中得到待删除数字的下标,将数组中该下标元素与数组末尾元素调换位置,再删除数组末尾元素,这样就做到了时间复杂度$O(1)$的删除。
class RandomizedSet {
private:
unordered_map<int, int> hash;
vector<int> nums;
public:
RandomizedSet() {}
bool insert(int val) {
if (hash.count(val)) return false; //哈希表中已存在该数值,不用再插入
hash[val] = nums.size(); //哈希表的键是数值,值是数值在数组中的位置
nums.push_back(val); //数值放入数组中
return true;
}
bool remove(int val) {
if (!hash.count(val)) return false; //哈希表中不存在该数值,无法删除
nums[hash[val]] = nums.back(); //将数组尾部的数值放到待删除数值的位置
hash[nums.back()] = hash[val]; //在哈希表中更新原数组尾部数值的位置
hash.erase(val); //从哈希表中删除该数值
nums.pop_back(); //从数组中删除该数值
return true;
}
int getRandom() {
return nums[rand() % nums.size()]; //随机返回集合内一数值
}
};