Design a map that supports all following operations in average O(1) time:
class RandomMap {
void put(int key, int value);
Integer get(int key);
Integer getRandom(); // this should return a random value with equal probability
void remove(int key);
}The question is similar to https://leetcode.com/problems/insert-delete-getrandom-o1 but in this case it's a map instead of a set which makes things slightly more complicated.
From xkcd:
My java solution
https://leetcode.com/playground/p3oRcCBK
public class RandomMap {
private final Map<Integer, Integer> keyToIndex = new HashMap<>();
private final List<Entry> entries = new ArrayList<>();
public void put(int key, int value) {
if (keyToIndex.containsKey(key)) {
Entry entry = entries.get(keyToIndex.get(key));
entry.value = value;
} else {
keyToIndex.put(key, entries.size());
entries.add(new Entry(key, value));
}
}
public Integer get(int key) {
Integer at = keyToIndex.get(key);
return at != null ? entries.get(at).value : null;
}
public Integer getRandom() {
if (entries.isEmpty()) return null;
int at = ThreadLocalRandom.current().nextInt(entries.size());
return entries.get(at).value;
}
public void remove(int key) {
if (entries.isEmpty() || !keyToIndex.containsKey(key)) return;
int curr = keyToIndex.get(key);
int last = entries.size() - 1;
if (curr != last) {
Collections.swap(entries, curr, last);
Entry entry = entries.get(curr);
keyToIndex.replace(entry.key, curr);
}
entries.remove(last); // remove last element O(1)
keyToIndex.remove(key);
}
private static class Entry {
int key;
int value;
public Entry(int key, int value) {
this.key = key;
this.value = value;
}
}
}Comments (7)
