Avoid Flood in The City

medium heap greedy binary search hash table

Problem

Every lake starts empty. On day i, if rains[i] > 0 that lake fills with water — and if it was already full, it floods. If rains[i] == 0 it is a dry day: you must pick exactly one lake and empty it. Build an answer array where rainy days are -1 and each dry day holds the lake you drained, choosing any schedule that avoids every flood. If no schedule works, return an empty array.

Inputrains = [1, 2, 0, 0, 2, 1]

Output[-1, -1, 2, 1, -1, -1]

Lakes 1 and 2 fill on days 0 and 1. Day 2 drains lake 2, day 3 drains lake 1, so the rains on days 4 and 5 fall on empty lakes — no flood.

rains (comma-separated, 0 = dry day)

import bisect
def avoidFlood(rains):
    ans = [-1] * len(rains)
    full = {}              # lake -> day it last became full
    dry = []               # sorted available dry-day indices
    for i, lake in enumerate(rains):
        if lake == 0:
            bisect.insort(dry, i)   # this day is free to dry a lake
            ans[i] = 1              # placeholder (overwritten if used)
        else:
            if lake in full:
                # need a dry day strictly after full[lake]
                pos = bisect.bisect_right(dry, full[lake])
                if pos == len(dry):
                    return []        # nothing free -> unavoidable flood
                ans[dry.pop(pos)] = lake
            full[lake] = i           # mark lake full on day i
    return ans

function avoidFlood(rains) {
  const ans = new Array(rains.length).fill(-1);
  const full = new Map();          // lake -> day it last became full
  const dry = [];                  // sorted available dry-day indices
  for (let i = 0; i < rains.length; i++) {
    const lake = rains[i];
    if (lake === 0) {
      dry.push(i);                 // appended in order, so dry stays sorted
      ans[i] = 1;                  // placeholder (overwritten if used)
    } else {
      if (full.has(lake)) {
        // first dry day strictly after full[lake]
        const pos = upperBound(dry, full.get(lake));
        if (pos === dry.length) return [];   // no free day -> flood
        ans[dry[pos]] = lake;
        dry.splice(pos, 1);
      }
      full.set(lake, i);           // mark lake full on day i
    }
  }
  return ans;
}

function upperBound(a, x) {        // first index with a[index] > x
  let lo = 0, hi = a.length;
  while (lo < hi) {
    const mid = (lo + hi) >> 1;
    if (a[mid] <= x) lo = mid + 1; else hi = mid;
  }
  return lo;
}

int[] avoidFlood(int[] rains) {
    int n = rains.length;
    int[] ans = new int[n];
    Arrays.fill(ans, -1);
    Map<Integer, Integer> full = new HashMap<>();  // lake -> last full day
    TreeSet<Integer> dry = new TreeSet<>();         // available dry days
    for (int i = 0; i < n; i++) {
        int lake = rains[i];
        if (lake == 0) {
            dry.add(i);                // this day is free to dry a lake
            ans[i] = 1;                // placeholder (overwritten if used)
        } else {
            if (full.containsKey(lake)) {
                // earliest dry day strictly after full[lake]
                Integer day = dry.higher(full.get(lake));
                if (day == null) return new int[0];  // no free day -> flood
                ans[day] = lake;
                dry.remove(day);
            }
            full.put(lake, i);         // mark lake full on day i
        }
    }
    return ans;
}

vector<int> avoidFlood(vector<int>& rains) {
    int n = rains.size();
    vector<int> ans(n, -1);
    unordered_map<int,int> full;     // lake -> day it last became full
    set<int> dry;                    // available dry-day indices
    for (int i = 0; i < n; i++) {
        int lake = rains[i];
        if (lake == 0) {
            dry.insert(i);          // this day is free to dry a lake
            ans[i] = 1;             // placeholder (overwritten if used)
        } else {
            auto it = full.find(lake);
            if (it != full.end()) {
                // first dry day strictly after the lake last filled
                auto d = dry.upper_bound(it->second);
                if (d == dry.end()) return {};  // no free day -> flood
                ans[*d] = lake;
                dry.erase(d);
            }
            full[lake] = i;         // mark lake full on day i
        }
    }
    return ans;
}

Explanation

This is a greedy scheduling problem. The only freedom we have is choosing which lake to drain on each dry day, and the only danger is a lake that fills twice with no drain in between.

Scan the days left to right. Keep full, a hash map from each currently-full lake to the day it filled, and dry, an ordered collection of the dry-day indices we have seen but not yet spent.

When day i is a dry day (rains[i] == 0) we just record it in dry and leave a placeholder in the answer; we postpone deciding what to drain until a lake actually needs saving.

When it rains on a lake that is already full, we must have drained it on some earlier dry day. The greedy rule: pick the earliest dry day that comes strictly after the day the lake filled. Using the soonest valid day keeps later dry days available for other lakes — never spending a day too early. A TreeSet / sorted set / heap lets us find and remove that day with higher / upper_bound (an O(log n) search), and we overwrite that day's answer with this lake.

If no dry day exists after the lake filled, the lake is guaranteed to flood, so we return an empty array. Otherwise every rainy day is -1, each spent dry day names the lake it saved, and any unused dry day keeps its harmless placeholder.

Example: rains = [1,2,0,0,2,1]. Lakes 1 and 2 fill on days 0 and 1. When lake 2 rains again on day 4 we grab dry day 2 (the earliest after day 1); when lake 1 rains again on day 5 we grab dry day 3. The result is [-1,-1,2,1,-1,-1].

Time: O(n log n) Space: O(n)