Problem Statement in English

You’re given an array rains where:

  • rains[i] > 0 means that on the i-th day, it rains over the lake with index rains[i].
  • rains[i] == 0 means that on the i-th day, you can choose one lake to dry.

You need to return an array ans where:

  • ans[i] == -1 if rains[i] > 0
  • ans[i] is the lake you choose to dry on the i-th day if rains[i] == 0

If there is no way to avoid a flood, return an empty array.

The question also says that emptying a lake that is already empty has no effect.

Approach

We can use a greedy approach with a priority queue (min-heap) to keep track of the lakes that are filled and the days they will next rain. The idea is to always dry the lake that will next rain the soonest when we have a 0 day.

Solution in Python


class Solution:
    def avoidFlood(self, rains: List[int]) -> List[int]:
        heap = []
        filled = set()
        nextRain = defaultdict(deque)
        ans = []

        for i in range(len(rains)):
            nextRain[rains[i]].append(i)

        for i, l in enumerate(rains):
            if l in filled:
                return []
            
            if not l:
                if heap:
                    _, fl = heappop(heap)
                    filled.remove(fl)
                    ans.append(fl)
                else:
                    ans.append(1)
            else:
                ans.append(-1)

                if nextRain[l] and nextRain[l][0] == i:
                    nextRain[l].popleft()

                if nextRain[l]:
                    ni = nextRain[l].popleft()
                    heappush(heap, (ni - i, l))
                    filled.add(l)
        
        return ans

Complexity

  • Time: $O(N \log N)$
    The time complexity is primarily determined by the operations on the priority queue (heap). In the worst case, we may perform a push and pop operation for each day, leading to a complexity of $O(N \log N)$.

  • Space: $O(N)$
    The space complexity is primarily determined by the storage used for the priority queue, the set of filled lakes, and the deque for the next rain days. In the worst case, we may store all lakes in these data structures.

And we are done.