Avoid Flood in The City LeetCode solution in Python

1
from bisect import bisect_left
2

3

4
class Solution:
5
    def avoidFlood(self, rains):
6
        full_lakes, dry_dates = {}, []
7
        ans = [-1] * len(rains)
8

9
        for date, rain_lake in enumerate(rains):
10
            if rain_lake == 0:  # no rain, we can dry one lake
11
                dry_dates.append(date)  # keep dry date, we'll decide later
12
                continue
13

14
            if rain_lake in full_lakes:  # the lake is already full
15
                # BS find out earliest day we can use to dry that lake | greedy
16
                dry_day = bisect_left(dry_dates, full_lakes[rain_lake])
17

18
                if dry_day >= len(dry_dates):
19
                    return []  # can not find a date to dry this lake
20

21
                ans[dry_dates.pop(dry_day)] = (
22
                    rain_lake  # dry this lake at the date we choose
23
                )
24

25
            # remember latest rain on this lake
26
            full_lakes[rain_lake] = date
27

28
        # we may have dry dates remain, on these days, rain > 0, we can not use -1, just choose day 1 to dry (maybe nothing happend)
29
        for dry_day in dry_dates:
30
            ans[dry_day] = 1
31

32
        return ans