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
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;
}
};