1488. 避免洪水泛滥

class Solution {
    public int[] avoidFlood(int[] rains) {
        int[] res = new int[rains.length];
        Arrays.fill(res, 1);
        TreeSet<Integer> freeDay = new TreeSet<>();
        Map<Integer, Integer> map = new HashMap<>();
        for (int i = 0; i < rains.length; i++) {
            int rain = rains[i];
            if (rain == 0) {
                // 空闲日子
                freeDay.add(i);
            } else {
                // 下雨
                res[i] = -1;
                if (map.containsKey(rain)) {
                    // 已经有水，则抽掉
                    Integer tp = freeDay.ceiling(map.get(rain));
                    if (tp == null) {
                        return new int[]{};
                    } else {
                        // 用之前存储的空闲天数抽掉 rain 湖泊的水
                        res[tp] = rain;
                        freeDay.remove(tp);
                    }
                }
                map.put(rain, i);
            }
        }
        return res;
    }
}