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();
The idea is to add a set to the hashMap to remember all the locations of a duplicated number.
1 public class RandomizedCollection { 2 HashMap<Integer, HashSet<Integer>> map; 3 ArrayList<Integer> arr; 4 java.util.Random random; 5 6 /** Initialize your data structure here. */ 7 public RandomizedCollection() { 8 map = new HashMap<Integer, HashSet<Integer>>(); 9 arr = new ArrayList<Integer>(); 10 random = new java.util.Random(); 11 } 12 13 /** Inserts a value to the collection. Returns true if the collection did not already contain the specified element. */ 14 public boolean insert(int val) { 15 boolean res = false; 16 if (!map.containsKey(val)) { 17 map.put(val, new HashSet<Integer>()); 18 res = true; 19 } 20 arr.add(val); 21 map.get(val).add(arr.size()-1); 22 return res; 23 } 24 25 /** Removes a value from the collection. Returns true if the collection contained the specified element. */ 26 public boolean remove(int val) { 27 if (!map.containsKey(val)) return false; 28 int lastItem = arr.get(arr.size()-1); 29 int index = arr.size()-1; 30 if (lastItem != val) { 31 HashSet<Integer> lastItemSet = map.get(lastItem); 32 index = map.get(val).iterator().next(); 33 arr.set(index, lastItem); 34 lastItemSet.remove(arr.size()-1); 35 lastItemSet.add(index); 36 } 37 38 if (map.get(val).size() == 1) map.remove(val); 39 else map.get(val).remove(index); 40 arr.remove(arr.size()-1); 41 return true; 42 } 43 44 /** Get a random element from the collection. */ 45 public int getRandom() { 46 return arr.get(random.nextInt(arr.size())); 47 } 48 } 49 50 /** 51 * Your RandomizedCollection object will be instantiated and called as such: 52 * RandomizedCollection obj = new RandomizedCollection(); 53 * boolean param_1 = obj.insert(val); 54 * boolean param_2 = obj.remove(val); 55 * int param_3 = obj.getRandom(); 56 */