LeetCode: Avoid Flood in The City (Midium)

問題はこちら

leetcode.com

湖が無限個存在している。その全ては今水は入っていない（=空）だが、 番目の湖に雨が降ると満杯になる。

満杯の状態の湖に雨が降ると、洪水が起きてしまう。

以下のような rains 配列が与えられるので、洪水を避けるようにしたい。

• rains[i] > 0: rains[i] 番目の湖に雨が降る。
• rains[i] == 0: 雨が降らず、好きな湖を一つ空にすることができる。

答えとしては、以下のようなans 配列を返す。

• ans.length == rains.length
• ans[i] == -1: rains[i] > 0 の時
• ans[i]: rains[i] == 0 の時、 日目に空にした湖の番号

考え方

空である湖 に雨が降った時は何も気にしなくて良い。しかし、次に湖 に再び雨が降った時、考えなければならない。よって、同じ湖 に雨が降った時に、それより前にある「空にする日」を用いて、その湖 を空にできなかったかを考える。そうすることで、空の湖に雨が降ることになるので洪水を避けることができる。

具体的には、次のような「空にする日」を選ぶのが良い。

• 前回湖に雨が降った日 prev_day よりも後。
  • 雨が降った日より後でないと、空にする意味がないため。
• かつ prev_day から一番近い日。
  • greedyに考える。 prev_day よりも遠いところにあるものほど、のちに使える可能性が高い。よって、なるべく使えるならば早めに使ってしまう方がよい。

実装方針

上記のような、「ある時点からもっとも近い時点を高速に求める」のは lower_boundを使えばできそうで、「空にする日」を set に入れることで探索を実現できる。

また、prev_day を求めるために、毎回雨が湖 に降った場合、その日付を記録しておく map を用意しておく。詳しくはコードを見た方が良いと思う。

解答

class Solution {
public:
    vector<int> avoidFlood(vector<int> &rains) {
        int n = rains.size();
        vector<int> v(n);

        // 「空にする日」を入れておく
        set<int> drydays;

        // full_lakes[i]: 今までで、湖iを水で満たした最後の日を持つ
        unordered_map<int, int> full_lakes;

        for (int i = 0; i < n; ++i) {
            if (rains[i] == 0) {
                drydays.insert(i);
                // 適当に入れる。後で書き換わる可能性あり
                v[i] = 1;
            } else {
                if (full_lakes.find(rains[i]) == full_lakes.end()) {
                    full_lakes[rains[i]] = i;
                } else {
                    if (drydays.empty()) return vector<int>();

                    // 前回に湖を満たした日を求める
                    int prev = full_lakes[rains[i]];

                    // prev以降でもっとも近いところに存在する、「空にする日」を探して、
                    // 空にする日はrains[i]を空にすることにする
                    auto itr = drydays.lower_bound(prev);
                    if (itr == drydays.end()) {
                        return vector<int>();
                    } else {
                        v[*itr] = rains[i];
                        drydays.erase(*itr);

                        // 湖を満たした日をアップデートする
                        full_lakes[rains[i]] = i;
                    }

                }
                v[i] = -1;
            }
        }

        return v;
    }
};