代码:
class RandomizedCollection381_hard {
public:
/** Initialize your data structure here. */
RandomizedCollection381_hard() {
}
/** Inserts a value to the collection. Returns true if the collection did not already contain the specified element. */
bool insert(int val) {
elem_list_.push_back(val);
elem_map_[val].insert(elem_list_.size() - 1);
return (elem_map_[val].size() == 1);
}
/** Removes a value from the collection. Returns true if the collection contained the specified element. */
bool remove(int val) {
if (elem_map_.find(val) == elem_map_.end())
return false;
int idx = *(elem_map_[val].begin());
elem_list_[idx] = elem_list_.back();
elem_map_[val].erase(idx);
/* 交换元素后需要更新被交换元素位置信息 */
elem_map_[elem_list_[idx]].erase(elem_list_.size() - 1);
/* 如果删除的是最末端元素,则直接删除 */
if (idx < elem_list_.size() - 1) {
elem_map_[elem_list_[idx]].insert(idx);
}
/* 在该元素结合为空的时候删除该元素的映射项 */
if (elem_map_[val].empty()) {
elem_map_.erase(val);
}
elem_list_.pop_back();
return true;
}
/** Get a random element from the collection. */
int getRandom() {
return elem_list_[rand() % elem_list_.size()];
}
private:
vector<int> elem_list_;
unordered_map<int, unordered_set<int>> elem_map_;
};