https://leetcode.com/problems/avoid-flood-in-the-city/

Source


    vector<int> avoidFlood(vector<int>& rains) {
        const int numDays = rains.size();
        vi viAns(numDays);
        unordered_map<int, int> hashLake_RainDay;
        set<int> setDryDays;
        FOR (day, numDays) {
            if (rains[day] == 0) {
                viAns[day] = 1; // default value
                setDryDays.emplace(day);
                continue;
            }

            viAns[day] = -1;
            const int fullOfWaterLake = rains[day];
            auto itLake_RainDay = hashLake_RainDay.find(fullOfWaterLake);
            if (end(hashLake_RainDay) == itLake_RainDay) {  // not found
                // first rain
                hashLake_RainDay[fullOfWaterLake] = day;
                continue;
            }

            // second rain
            if (setDryDays.empty()) {
                // inevitable flood
                return vi();
            }

            const int prevRainDay = itLake_RainDay->second;
            int dryDay = -1;
            auto itDryDay=begin(setDryDays);
#ifdef BINARY_SEARCH
            if (*itDryDay > prevRainDay) {
                // pick the first dry day
                dryDay = *itDryDay;
            }
            else {
                const int numDryDays = setDryDays.size();
                const int MAX_L = log2(numDryDays)+3;
                int hi = numDryDays-1;
                int lo = 0;
                for (int i=0; (hi>=lo) && i<MAX_L; ++i) {
                    const int mid = (hi+lo) >> 1;
                    // auto midVal = getNthElem(setDryDays, mid);
                    auto itMid = next(begin(setDryDays), mid);
                    if (*itMid >= prevRainDay) {
                        // dryDay = midVal.first;
                        dryDay = *itMid; itDryDay = itMid;
                        if (hi == mid) hi = mid-1;
                        else hi = mid;
                    }
                    else {
                        if (lo == mid) lo = mid+1;
                        else lo = mid;
                    }
                }
            }
#else
            for (; itDryDay!=end(setDryDays); ++itDryDay) {
                dryDay = *itDryDay;
                if (dryDay <= prevRainDay) {
                    // invalid dry day (too early)
                    continue;
                }
                break;
            }
            if (end(setDryDays) == itDryDay) {  // not found
                // inevitable flood
                return vi();
            }
#endif
            if (dryDay < 0) return vi();

            // consume a day to be used to dry a lake
            setDryDays.erase(itDryDay);
            // dry the lake at the valid dry day
            viAns[dryDay] = fullOfWaterLake;
            // dry a lake
            hashLake_RainDay.erase(itLake_RainDay);

            // but, it rains again...
            hashLake_RainDay[fullOfWaterLake] = day;
        }
        return viAns;
    }

    template <typename T>
    pair<T, bool> getNthElem(set<T>& searchSet, int n) {
        pair<T, bool> result;
        if (searchSet.size() > static_cast<size_t>(n)) {
            result.first = *(next(begin(searchSet), n));
            result.second = true;
        }
        else {
            result.second = false;
        }
        return result;
    }

GitHub

AvoidFlood