设计一个支持在平均 时间复杂度 O(1)
下, 执行以下操作的数据结构。
注意: 允许出现重复元素。
insert(val):向集合中插入元素 val。
remove(val):当 val 存在时,从集合中移除一个 val。
getRandom:从现有集合中随机获取一个元素。每个元素被返回的概率应该与其在集合中的数量呈线性相关。
示例:
// 初始化一个空的集合。
RandomizedCollection collection = new RandomizedCollection();
// 向集合中插入 1 。返回 true 表示集合不包含 1 。
collection.insert(1);
// 向集合中插入另一个 1 。返回 false 表示集合包含 1 。集合现在包含 [1,1] 。
collection.insert(1);
// 向集合中插入 2 ,返回 true 。集合现在包含 [1,1,2] 。
collection.insert(2);
// getRandom 应当有 2/3 的概率返回 1 ,1/3 的概率返回 2 。
collection.getRandom();
// 从集合中删除 1 ,返回 true 。集合现在包含 [1,2] 。
collection.remove(1);
// getRandom 应有相同概率返回 1 和 2 。
collection.getRandom();
解答
由于需要o(1)
复杂度插入,尤其是去除,因此可以使用哈希字典将数字和其位置索引进行映射:
class RandomizedCollection {
vector<int> nums;
unordered_map<int, unordered_set<int>> m;
public:
/** Initialize your data structure here. */
RandomizedCollection() {
srand((unsigned) time(0));
}
/** Inserts a value to the collection. Returns true if the collection did not already contain the specified element. */
bool insert(int val) {
auto iter = m.find(val);
bool insert_result = (iter == m.end() || iter->second.empty());
if(iter == m.end()){
unordered_set<int> temp;
m.emplace(make_pair(val, temp));
iter = m.find(val);
}
iter->second.insert(nums.size());
nums.emplace_back(val);
return insert_result;
}
/** Removes a value from the collection. Returns true if the collection contained the specified element. */
bool remove(int val) {
auto iter = m.find(val);
if(iter == m.end() || iter->second.empty())
return false;
// 去除待删除数字的索引
int target_index = *(iter->second.begin());
iter->second.erase(target_index);
// 如果该元素正好为数组中的最后一个,直接在数组中去除末尾即可
if(target_index == nums.size() - 1){
nums.pop_back();
}
// 否则,和数组末尾元素交换后再去除
else{
int last_num = nums.back();
swap(nums[nums.size() - 1], nums[target_index]);
auto iter_last = m.find(last_num);
iter_last->second.erase(nums.size()-1);
iter_last->second.insert(target_index);
nums.pop_back();
}
return true;
}
/** Get a random element from the collection. */
int getRandom() {
int random_idx = rand() % nums.size();
return nums[random_idx];
}
};
一开始本来是使用unordered_map<int, vector<int>>
来存储数字-位置索引的,但是vector
的插入和查找并非o(1)
复杂度,所以换成了unordered_set