1488. Avoid Flood in The City

Avoid Flood in The City - LeetCodeLeetCode

Weekly Contest 194 - Q3, 其实都不算是很难的Q3, 但就是一直TLE. 对于很多时间复杂度比较大的操作不够谨慎。思路也不够开阔，自己算法工具箱里可以拿出来用的东西太少了。

比赛时候的思路是维护所有lake的状态，如果遇到雨下在已经 full 的 lake 上，就往回找最早可用的 work day，把它给dry了。但按照这个思路写出来的代码一直 TLE （也附在下面了），过不了。

其实这个思路是可以做的，时间复杂度大概在O(n*logn)级别，但我的问题出在「找最早可用的work day」这一部分，我的操作是：

1. 先计算可用区间，

2. 再看work里是否存在可用的

这样需要对 work 搜索很多次（更傻的是，用remove删元素时还要再多搜一次），导致运算成本一下子上去了。其实只要把这边改成：遍历work，看有没有元素在可用区间内，即可通过。有效代码如下：

Greedy Straight Forward 解法

TLE Code

class Solution:
    def avoidFlood(self, rains: List[int]) -> List[int]:
        ans = [-1 for _ in range(len(rains))]
        full = collections.defaultdict(lambda: -1)
        work = set()
        for day, lake in enumerate(rains):
            if lake == 0:
                work.add(day)
            if lake > 0:
                if full[lake] < 0:
                    full[lake] = day
                else:
                    prev = full[lake]
                    if day - prev <= 1:
                        return []
                    found = False
                    for i in range(prev, day):
                        if i in work:
                            work.remove(i)
                            ans[i] = lake
                            full[lake] = day
                            found = True
                            break
                    if not found:
                        return []
        
        for i in work:
            ans[i] = 1
        
        return ans

有效提交 Code

class Solution:
    def avoidFlood(self, rains: List[int]) -> List[int]:
        work = []
        full = collections.defaultdict(lambda: -1)
        ans = []
        
        for day, lake in enumerate(rains):
            if lake == 0:
                work.append(day)
                ans.append(1)
                
            else:
                if full[lake] < 0:
                    full[lake] = day
                else:
                    prev = full[lake]
                    found = False
                    for i, d in enumerate(work):
                        if d > prev:
                            ans[d] = lake
                            full[lake] = day
                            del work[i]
                            found = True
                            break
                    if not found:
                        return []
                    
                ans.append(-1)
        
        return ans

Discussion里另一个常见解法就是利用Heap 维护一个危机序列 toDry， 一直处理危机序列里面最紧迫的危机（最小的 day index），本质也是一种 greedy 搜索。

Heap 解法

class Solution:
    def avoidFlood(self, rains: List[int]) -> List[int]:
        dayRain = collections.defaultdict(list)
        ans = [-1] * len(rains)
        toDry = []
        
        # 遍历 rains, 记录各个湖下雨的 day index
        for day,lake in enumerate(rains):
            dayRain[lake].append(day)
        
        for i, lake in enumerate(rains):
            # 如果下雨
            if lake != 0:
                # 危机可以解决 -> 添加危机
                if dayRain[lake] and len(dayRain[lake])>1:
                    heapq.heappush(toDry, dayRain[lake][1])
                # 危机不可能解决 -> Game over
                if dayRain[lake] and dayRain[lake][0] < i: return []
            # 如果不下雨 -> 可以干活
            else:
                # 有活儿可干 -> 解决最紧迫的危机 (toDry 里面最小的 day index)
                if toDry:
                    ans[i] = rains[heapq.heappop(toDry)]
                    dayRain[ans[i]].pop(0) # 活儿干完，当前危机解除
                # 没活儿干 -> 随便 dry 一个
                else: ans[i] = 1
        return ans