1488. 避免洪水泛滥 O(n log n) & O(n)【有序集合(插入,删除,有序)、二分查找】

32 阅读1分钟

1488. 避免洪水泛滥

from sortedcontainers import SortedList

class Solution:
    def avoidFlood(self, rains: List[int]) -> List[int]:
        # 有序集合 存储 没有下雨的日子  
        # 剩余晴天 用于抽干任意湖泊  令其为 1
        #  在 有序集合 st 中找  小于此次索引的 最小索引  二分查找
        res = [1] * len(rains)
        st = SortedList()  # 有序集合  记录晴天0的索引 
        dic = {} # 记录 每个湖泊 的上一次下雨的 索引
        for i, num in enumerate(rains):
            if num == 0:
                st.add(i)
            else:
                res[i] = - 1
                # 讨论能否 避免洪水
                if num in dic:
                    idx = st.bisect(dic[num]) # bisect.bisect(a,x)(默认等同于bisect.bisect_right())
                    if idx == len(st): # 没有 可用的晴天
                        return []
                    res[st[idx]] = num 
                    st.discard(st[idx])  # 删掉 该元素
                dic[num] = i # 更新该湖泊 下雨的索引

        return res 

image.png

class Solution {
public:
    vector<int> avoidFlood(vector<int>& rains) {
        vector<int> res(rains.size(), 1);
        set<int> st; // 存储 晴天的索引
        unordered_map<int, int> dic;
        for (int i = 0; i < rains.size(); ++i){
            if (rains[i] == 0){
                st.insert(i);
            }else{
                res[i] = -1;
                if (dic.count(rains[i])){
                    auto it = st.lower_bound(dic[rains[i]]); // 大于等于 上一次下雨的 索引
                    if (it == st.end()){
                        return {};
                    }
                    res[*it] = rains[i];
                    st.erase(it);
                }
                dic[rains[i]] = i; // 更新 该湖泊下雨的索引
            }
        }
        return res;
    }
};