Skip to content

1488. Avoid Flood in The City 👍

  • Time: $O(n\log n)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
class Solution {
 public:
  vector<int> avoidFlood(vector<int>& rains) {
    vector<int> ans(rains.size(), -1);
    unordered_map<int, int> lakeIdToFullDay;
    set<int> emptyDays;  // indices of rains[i] == 0

    for (int i = 0; i < rains.size(); ++i) {
      const int lakeId = rains[i];
      if (lakeId == 0) {
        emptyDays.insert(i);
        continue;
      }
      if (const auto itFullDay = lakeIdToFullDay.find(lakeId);
          itFullDay != lakeIdToFullDay.cend()) {
        // The lake was full in a previous day. Greedily find the closest day
        // to make the lake empty.
        const auto itEmptyDay = emptyDays.upper_bound(itFullDay->second);
        if (itEmptyDay == emptyDays.cend())  // Not found.
          return {};
        // Empty the lake at this day.
        ans[*itEmptyDay] = lakeId;
        emptyDays.erase(itEmptyDay);
      }
      // The lake with `lakeId` becomes full at the day `i`.
      lakeIdToFullDay[lakeId] = i;
    }

    // Empty an arbitrary lake if there are remaining empty days.
    for (const int emptyDay : emptyDays)
      ans[emptyDay] = 1;

    return ans;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
class Solution {
  public int[] avoidFlood(int[] rains) {
    int[] ans = new int[rains.length];
    Arrays.fill(ans, -1);
    Map<Integer, Integer> lakeIdToFullDay = new HashMap<>();
    TreeSet<Integer> emptyDays = new TreeSet<>(); // indices of rains[i] == 0

    for (int i = 0; i < rains.length; ++i) {
      final int lakeId = rains[i];
      if (lakeId == 0) {
        emptyDays.add(i);
        continue;
      }
      if (lakeIdToFullDay.containsKey(lakeId)) {
        final int fullDay = lakeIdToFullDay.get(lakeId);
        // The lake was full in a previous day. Greedily find the closest day
        // to make the lake empty.
        Integer emptyDay = emptyDays.higher(fullDay);
        if (emptyDay == null) // Not found.
          return new int[] {};
        // Empty the lake at this day.
        ans[emptyDay] = lakeId;
        emptyDays.remove(emptyDay);
      }
      // The lake with `lakeId` becomes full at the day `i`.
      lakeIdToFullDay.put(lakeId, i);
    }

    // Empty an arbitrary lake if there are remaining empty days.
    for (final int emptyDay : emptyDays)
      ans[emptyDay] = 1;

    return ans;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
from sortedcontainers import SortedSet


class Solution:
  def avoidFlood(self, rains: list[int]) -> list[int]:
    ans = [-1] * len(rains)
    lakeIdToFullDay = {}
    emptyDays = SortedSet()  # indices of rains[i] == 0

    for i, lakeId in enumerate(rains):
      if lakeId == 0:
        emptyDays.add(i)
        continue
        # The lake was full in a previous day. Greedily find the closest day
        # to make the lake empty.
      if lakeId in lakeIdToFullDay:
        fullDay = lakeIdToFullDay[lakeId]
        emptyDayIndex = emptyDays.bisect_right(fullDay)
        if emptyDayIndex == len(emptyDays):  # Not found.
          return []
        # Empty the lake at this day.
        emptyDay = emptyDays[emptyDayIndex]
        ans[emptyDay] = lakeId
        emptyDays.discard(emptyDay)
      # The lake with `lakeId` becomes full at the day `i`.
      lakeIdToFullDay[lakeId] = i

    # Empty an arbitrary lake if there are remaining empty days.
    for emptyDay in emptyDays:
      ans[emptyDay] = 1

    return ans