From:
Design a data structure that offers the following operations in O(1) time:
- insert
- remove
- contains
- get random element
Solution:
Consider a data structure composed of a hashtable H and an array A. The hashtable keys are the elements in the data structure, and the values are their positions in the array.
- insert(value): append the value to array and let i be it's index in A. Set H[value]=i.
- remove(value): We are going to replace the cell that contains value in A with the last element in A. let d be the last element in the array A at index m. let i be H[value], the index in the array of the value to be removed. Set A[i]=d, H[d]=i, decrease the size of the array by one, and remove value from H.
- contains(value): return H.contains(value)
- getRandomElement(): let r=random(current size of A). return A[r].
since the array needs to auto-increase in size, it's going to be amortize O(1) to add an element, but I guess that's OK.
public class ConstantRandomGet<E> {
private Map<E, Integer> map = new HashMap<>();
private List<E> list = new ArrayList<>();
private Random random = new Random();
public void insert(E e) {
map.put(e, list.size());
list.add(e);
}
public boolean remove(E e) {
if(!map.containsKey(e)) return false;
int i = map.get(e);
int last = list.size()-1;
E lastE = list.remove(last);
map.put(lastE, i);
list.set(i, lastE);
return true;
}
public boolean contains(E e) {
return map.containsKey(e);
}
public E getRandomElement() {
int size = list.size();
if(size == 0) return null;
int index = random.nextInt(size);
return list.get(index);
}
}