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();
Solution1: HashMap + ArrayList
code
1 public class RandomizedCollection { 2 class Node { 3 public int value; 4 public int index; 5 public Node(int val, int idx) { 6 value = val; 7 index = idx; 8 } 9 } 10 11 private Map<Integer, List<Integer>> map; 12 private List<Node> list; 13 private Random r; 14 15 /** Initialize your data structure here. */ 16 public RandomizedCollection() { 17 map = new HashMap<>(); 18 list = new ArrayList<>(); 19 r = new Random(); 20 } 21 22 /** Inserts a value to the collection. Returns true if the collection did not already contain the specified element. */ 23 public boolean insert(int val) { 24 List<Integer> l = map.getOrDefault(val, new ArrayList<>()); 25 l.add(list.size()); 26 map.put(val, l); 27 list.add(new Node(val, l.size() - 1)); 28 return l.size() == 1; 29 } 30 31 /** Removes a value from the collection. Returns true if the collection contained the specified element. */ 32 public boolean remove(int val) { 33 if (!map.containsKey(val)) return false; 34 List<Integer> l = map.get(val); 35 int removeIdx = l.get(l.size() - 1); 36 Node replaceNode = list.get(list.size() - 1); 37 38 // deal with HashMap 39 map.get(replaceNode.value).set(replaceNode.index, removeIdx); 40 l.remove(l.size() - 1); 41 if (l.size() == 0) map.remove(val); 42 43 // deal with List 44 list.set(removeIdx, replaceNode); 45 list.remove(list.size() - 1); 46 47 return true; 48 } 49 50 /** Get a random element from the collection. */ 51 public int getRandom() { 52 return list.get(r.nextInt(list.size())).value; 53 } 54 }