Avoid Flood in The City

Your country has an infinite number of lakes. Initially, all the lakes are empty, but when it rains over the nth lake, the nth the 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 the flood in any lake.

Given an integer array rains where:

1.rains[i] > 0 means there will be rains over the rains[i] lake.

2.rains[i] == 0 means there are no rains this day and you can choose one lake this day and dry it.

Return an array ans where:

1. ans.length == rains.length

2. ans[i] == -1 if rains[i] > 0.

3.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.

 Example:

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.

Approach:

C++

#include <bits/stdc++.h>
using namespace std;

vector<int> avoidFlood(vector<int> &rains)
{
    set<int> sun;
    unordered_map<int, int> last_rain;
    vector<int> ans(rains.size(), -1);
    for (int i = 0; i < rains.size(); ++i)
    {
        if (rains[i] == 0)
        {
            sun.insert(i);
            ans[i] = 1;
        }
        else
        {
            if (last_rain.count(rains[i]) > 0)
            {
                // the lake is full, we should
                // find the first sun day after the last rain day
                auto it = sun.upper_bound(last_rain[rains[i]]);
                // can not find the way to dry the lake
                if (it == end(sun))
                {
                    return {};
                }
                ans[*it] = rains[i];
                sun.erase(it);
            }
            // mark the last rain day
            last_rain[rains[i]] = i;
        }
    }
    return ans;
}

int main()
{
    vector<int> rains = {1, 2, 0, 0, 2, 1};

    vector<int> res = avoidFlood(rains);

    cout << "[";
    for (int i = 0; i < res.size(); i++)
    {
        cout << res[i];
        if (i != res.size() - 1)
            cout << ",";
    }
    cout << "]";

    return 0;
}