题目:
Design a data structure that supports all following operations in average O(1) time.
Note: Duplicate elements are allowed.
insert(val)
: Inserts an item val to the collection.remove(val)
: Removes an item val from the collection if present.getRandom
: Returns a random element from current collection of elements. The probability of each element being returned is linearly related to the number of same value the collection contains.
Example:
// Init an empty collection. RandomizedCollection collection = new RandomizedCollection(); // Inserts 1 to the collection. Returns true as the collection did not contain 1. collection.insert(1); // Inserts another 1 to the collection. Returns false as the collection contained 1. Collection now contains [1,1]. collection.insert(1); // Inserts 2 to the collection, returns true. Collection now contains [1,1,2]. collection.insert(2); // getRandom should return 1 with the probability 2/3, and returns 2 with the probability 1/3. collection.getRandom(); // Removes 1 from the collection, returns true. Collection now contains [1,2]. collection.remove(1); // getRandom should return 1 and 2 both equally likely. collection.getRandom();
思路:
有了上一道题目的基础,这道题目就好办了:思路唯一不同的是:对于哈希表中的每个哈希值,它可能会对应多个索引,因此将哈希表的value再设置成一个哈希表。
代码:
class RandomizedCollection {
public:
/** Initialize your data structure here. */
RandomizedCollection() {
}
/** Inserts a value to the collection. Returns true if the collection did not already contain the specified element. */
bool insert(int val) {
bool ret = (hash.count(val) == 0);
hash[val].insert(values.size());
values.push_back(val);
return ret;
}
/** Removes a value from the collection. Returns true if the collection contained the specified element. */
bool remove(int val) {
if(hash.count(val) == 0) {
return false;
}
else {
int last_value = values[values.size() - 1]; // the values of in the last
int last_index = values.size() - 1; // the index of the last
if(last_value == val) {
hash[val].erase(last_index);
if(hash[val].size() == 0) {
hash.erase(val);
}
values.pop_back();
}
else {
int index = *(hash[val].begin()); // the index to be modified
hash[last_value].erase(last_index);
hash[last_value].insert(index);
hash[val].erase(index);
if(hash[val].size() == 0) {
hash.erase(val);
}
values[index] = last_value;
values.pop_back();
}
return true;
}
}
/** Get a random element from the collection. */
int getRandom() {
if(values.size() == 0) {
return -1;
}
int rand_index = rand() % values.size();
return values[rand_index];
}
private:
unordered_map<int, unordered_set<int>> hash; // map from val to indices
vector<int> values; // map from index to val
};
/**
* Your RandomizedCollection object will be instantiated and called as such:
* RandomizedCollection obj = new RandomizedCollection();
* bool param_1 = obj.insert(val);
* bool param_2 = obj.remove(val);
* int param_3 = obj.getRandom();
*/