855. Exam Room!!

class ExamRoom {

    Map<Integer, int[]> startMap;
    Map<Integer, int[]> endMap;
    TreeSet<int[]> treeSet;
    int size;
    public ExamRoom(int n) {
        startMap = new HashMap<>();
        endMap = new HashMap<>();
        treeSet = new TreeSet<>((a,b)->{
            int dA=distance(a);
            int dB=distance(b);
            if(dA==dB){
                return b[0]-a[0];
            }else{
                return dA-dB;
            }
        });
        size=n;
        // seat -1, n will never be reached
        add(new int[]{-1, n});
    }
    
    public int seat() {
        // get the largest interval from treeSet;
        int[] interval = treeSet.last();
        // represent the position of the new seat;
        int index = 0;
        if(interval[0]==-1){
            index=0;
        }else if(interval[1]==size){
            index = size-1;
        }else{
            index = interval[0] + (interval[1]-interval[0])/2;
        }
        
        remove(interval);
        add(new int[]{interval[0],index});
        add(new int[]{index, interval[1]});
        return index;
    }
    
    public void leave(int p) {
        int[] interval = new int[]{endMap.get(p)[0],startMap.get(p)[1]};
        remove(startMap.get(p));
        remove(endMap.get(p));
        add(interval);
    }
    
    private void add(int[] interval){
        treeSet.add(interval);
        startMap.put(interval[0], interval);
        endMap.put(interval[1], interval);
    }
    
    private void remove(int[] interval){
        // remove: remove from tree set;
        treeSet.remove(interval);
        // remember: remove a element from HashMap
        startMap.remove(interval[0]);
        endMap.remove(interval[1]);
    }
    
    private int distance(int[] interval){
        // if interval[0]==-1, means the seat should be at index 0,
        // so that distance to end can be maximized;
        if(interval[0]==-1) return interval[1];
        
        if(interval[1]==size) return interval[1]-1-interval[0];
        return (interval[1]-interval[0])/2;
    }
}

/**
 * Your ExamRoom object will be instantiated and called as such:
 * ExamRoom obj = new ExamRoom(n);
 * int param_1 = obj.seat();
 * obj.leave(p);
 */

Last updated

Was this helpful?