528. Random Pick with Weight

Link

Solution

Record the range for each index by treeMap.

And random a number by the total range and use treeMap to find the correct one.

   TreeMap<Integer , Integer> tmap;
    Random random;
    int total = 0;
    public Solution(int[] w) {
        tmap = new TreeMap<>();
        int i = 0;
        for(int num : w){
            total += num;
            tmap.put(total-1, i++); //if the random num == total, then it is not the case. so we should remove it.
        }
        random = new Random();
    }
    
    public int pickIndex() {
        int r = random.nextInt(total); //random from 0 ~ total-1
        return tmap.ceilingEntry(r).getValue();
    }