leetcode 1488 避免洪水泛滥

用TreeSet实现二分法的查找

class Solution {
    public int[] avoidFlood(int[] rains) {
        int[] ans = new int[rains.length];
        Arrays.fill(ans, 1);
        // TreeSet自动实现二分法排序
        TreeSet<Integer> st = new TreeSet<Integer>();
        //存储当前池子上一次下雨的时间，如果后面满了，找到上一次满的时间之后的第一个晴天日期
        Map<Integer, Integer> mp = new HashMap<Integer, Integer>();
        for (int i = 0; i < rains.length; ++i) {
            if (rains[i] == 0) {
                st.add(i);
            } else {
                ans[i] = -1;
                if (mp.containsKey(rains[i])) {
                    Integer it = st.ceiling(mp.get(rains[i]));
                    if (it == null) {
                        return new int[0];
                    }
                    ans[it] = rains[i];
                    st.remove(it);
                }
                mp.put(rains[i], i);
            }
        }
        return ans;
    }
}