Java O(1) solution with correct getRandom() realisation


  • 0
    W

    public class RandomizedCollection {
    // Map to keep all numbers with an index assigned when inserted into this collection.
    private Map<Integer, Integer> idxNum;
    // Map to keep all the indexes assigned in idxNum of a number in this collection.
    private Map<Integer, Set<Integer>> numLocs;
    // Number of elements in this collection.
    private int count = 0;
    Random r;

    /** Initialize your data structure here. */
    public RandomizedCollection() {
        idxNum = new HashMap<>();
        numLocs = new HashMap<>();
        r = new Random();
    }
    
    /** Inserts a value to the collection. Returns true if the collection did not already contain the specified element. */
    public boolean insert(int val) {
        // Put val into collection ans assign index (count+1) to it.
        count ++;
        idxNum.put(count, val);
    
        // Add this new index of val to numLocs.
        Set<Integer> locs = numLocs.get(val);
        boolean isExist = (locs != null);
        
        if(!isExist) {
            locs = new HashSet<Integer>();
        }
        locs.add(count);
        numLocs.put(val, locs);
        
        return !isExist;
    }
    
    /** Removes a value from the collection. Returns true if the collection contained the specified element. */
    public boolean remove(int val) {
        Set<Integer> locs = numLocs.get(val);      
        if(locs == null) return false;
    
        Iterator<Integer> itr = locs.iterator();
        
        // Remove the element with index itr.next() from location map.
        int idx = itr.next();
        itr.remove();
        if(locs.size() == 0) {
            numLocs.remove(val);
        } else {
            numLocs.put(val, locs);
        }
        
        // Update index map idxNum by swapping element with index idx and count and removing the max index element.
        if(idx != count) {
            int maxIdxEle = idxNum.get(count);
            idxNum.put(idx, maxIdxEle);
    
            // Update locations of maxIdxEle
            Set<Integer> locsOfMaxIdxEle =  numLocs.get(maxIdxEle);
            locsOfMaxIdxEle.remove(count);
            locsOfMaxIdxEle.add(idx);
            numLocs.put(maxIdxEle, locsOfMaxIdxEle);            
        }
        idxNum.remove(count);
        count --;
    
        return true;
    }
    
    /** Get a random element from the collection. */
    public int getRandom() {
        if(count == 0) return -1;
        int index = r.nextInt(count);
        return idxNum.get(index+1);
    }
    

    }


Log in to reply
 

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