1488. Avoid Flood in The City

class Solution {
    public int[] avoidFlood(int[] rains) {
        int[] ans = new int[rains.length];
        Arrays.fill(ans, 1);

        TreeSet<Integer> sunDaySet = new TreeSet<>();
        Map<Integer, Integer> map = new HashMap<>();

        for (int i = 0; i < rains.length; i++) {
            if (rains[i] > 0) {
                // 第 rains[i] 个湖泊会下雨
                ans[i] = -1; // 题目要求返回 -1

                // 判断是否为重复下雨，因为抽干会重复下雨的湖泊才能最大化利用不下雨的天气
                if (map.containsKey(rains[i])) {
                    // rains[i] 重复下雨，看是否有晴天能用来抵消
                    Integer sunDay = sunDaySet.higher(map.get(rains[i])); // 注意此处是 map.get(rains[i]), 因为是寻找该湖泊上次下雨后更大的最近的天数
                    if (sunDay == null) {
                        // 没有晴天可以抵消
                        return new int[0];
                    } else {
                        ans[sunDay] = rains[i];
                        sunDaySet.remove(sunDay);
                    }
                }
                map.put(rains[i], i); // 记录下该湖泊最近下雨的天数
            } else {
                // 第 i 天没有湖泊会下雨，将 i 存下来，当后面重复下雨时用来抵消
                sunDaySet.add(i);
            }
        }

        return ans;
    }
}

References

1488. Avoid Flood in The City