Description (leetcode)

Your country has 10^9 lakes. Initially, all the lakes are empty, but when it rains over the n^th lake, the n^th lake becomes full of water. If it rains over a lake that is full of water, there will be a flood. Your goal is to avoid floods in any lake.

Given an integer array rains where:

  • rains[i] > 0 means there will be rains over the rains[i] lake.
  • rains[i] == 0 means there are no rains this day and you must choose one lake this day and dry it.

Return an array ans where:

  • ans.length == rains.length
  • ans[i] == -1 if rains[i] > 0.
  • ans[i] is the lake you choose to dry in the ith day if rains[i] == 0.

If there are multiple valid answers return any of them. If it is impossible to avoid flood return an empty array.

Notice that if you chose to dry a full lake, it becomes empty, but if you chose to dry an empty lake, nothing changes.

Example 1:

Input: rains = [1,2,3,4]

Output: [-1,-1,-1,-1]

Explanation: After the first day full lakes are [1]

After the second day full lakes are [1,2]

After the third day full lakes are [1,2,3]

After the fourth day full lakes are [1,2,3,4]

There’s no day to dry any lake and there is no flood in any lake.

Example 2:

Input: rains = [1,2,0,0,2,1]

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

Explanation: After the first day full lakes are [1]

After the second day full lakes are [1,2]

After the third day, we dry lake 2. Full lakes are [1]

After the fourth day, we dry lake 1. There is no full lakes.

After the fifth day, full lakes are [2].

After the sixth day, full lakes are [1,2].

It is easy that this scenario is flood-free. [-1,-1,1,2,-1,-1] is another acceptable scenario.

Example 3:

Input: rains = [1,2,0,1,2]

Output: []

Explanation: After the second day, full lakes are [1,2]. We have to dry one lake in the third day.

After that, it will rain over lakes [1,2]. It’s easy to prove that no matter which lake you choose to dry in the 3rd day, the other one will flood.

Constraints:

  • 1 <= rains.length <= 10^5
  • 0 <= rains[i] <= 10^9

submission

use std::collections::HashMap;
use std::collections::BTreeSet;

impl Solution {
    // Intuition: We delay the decision of pick a lake to dry
    // to the day we have to avoid flood
    // keep track the flooded lake, keep a set of days
    // are allowed to dry a lake, iterate through days, when we
    // encoutered a flooded lake in raining day, we take a day from
    // the set
    pub fn avoid_flood(rains: Vec<i32>) -> Vec<i32> {
        // keep track the states of lakes
        // the value is the previous rain day
        let mut flooded = HashMap::new();
        // set of days of no rains
        let mut days = BTreeSet::new();
        let mut schedule = vec![-1; rains.len()];
        for (idx, &lake) in rains.iter().enumerate() {
            match lake {
                0 => {
                    // add this day to set
                    days.insert(idx);
                    // leetcode will not accept our answer if
                    // we return -1 when not raining
                    schedule[idx] = 1;
                }
                lake => {
                    if let Some(filled) = flooded.get(&lake) {
                        // we need to schedule a way to dry this lake
                        let Some(&day) = days.range(filled..).next() else {
                            // we don't have any day available
                            return vec![];
                        };
                        days.remove(&day);
                        schedule[day] = lake;
                    }
                    // update flooded state
                    flooded.insert(lake, idx);
                }
            }
        }
        schedule
    }
}