Java AC Solution(147ms), linearly related to the number of same value, but not exactly averaged O(1)


  • 0
    E
    public class RandomizedCollection {
        ArrayList<Integer> nums;
        HashMap<Integer, PriorityQueue<Integer>> locmap;
        java.util.Random rand = new java.util.Random();
        
        /** Initialize your data structure here. */
        public RandomizedCollection() {
            nums = new ArrayList<Integer>();
            locmap = new HashMap<Integer, PriorityQueue<Integer>>();
        }
        
        /** Inserts a value to the collection. Returns true if the collection did not already contain the specified element. */
        public boolean insert(int val) {
            // biggeest position is always at the top of heap
            PriorityQueue<Integer> l = new PriorityQueue<Integer>(new Comparator<Integer>(){
                @Override
                public int compare(Integer o1, Integer o2){
                    return o2 - o1;
                }
            });
            if(locmap.containsKey(val)) l = locmap.get(val);
            l.add(nums.size()); // O(logk)
            locmap.put(val, l);
            nums.add(val);
            return true;
        }
        
        /** Removes a value from the set. Returns true if the collection contained the specified element. */
        public boolean remove(int val) {
            PriorityQueue<Integer> locq = locmap.get(val);
            if (locq == null || locq.size() == 0) return false;
            int endpos = locq.poll();
            // if not the last position, swap the last position with this position
            if (endpos < nums.size() - 1) { 
                int lastone = nums.get(nums.size()-1);
                PriorityQueue<Integer> tempq = locmap.get(lastone); 
                tempq.poll();
                tempq.add(endpos);  // O(logk)
                nums.set(endpos , lastone);
            }
            nums.remove(nums.size() - 1);
            return true;
        }
        
        /** Get a random element from the collection. */
        public int getRandom() {
            return nums.get(rand.nextInt(nums.size()));
        }
    }
    
    /**
     * Your RandomizedCollection object will be instantiated and called as such:
     * RandomizedCollection obj = new RandomizedCollection();
     * boolean param_1 = obj.insert(val);
     * boolean param_2 = obj.remove(val);
     * int param_3 = obj.getRandom();
     */
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.