utilities/test/predecessor_problem.test.cpp

• View this file on GitHub
• Last update: 2021-04-25 18:25:53+09:00
• Problem: https://judge.yosupo.jp/problem/predecessor_problem

Depends on

• utilities/integer_segments.hpp

Code

#include "../integer_segments.hpp"
#include <iostream>
#include <string>
using namespace std;
#define PROBLEM "https://judge.yosupo.jp/problem/predecessor_problem"

int main() {
    cin.tie(nullptr), ios::sync_with_stdio(false);

    int N, Q;
    string T;
    cin >> N >> Q >> T;
    integer_segments<int> seg;
    T += '0';
    int l = 0;
    for (int i = 0; i < int(T.size()); i++) {
        if (T[i] == '0') {
            if (l <= i - 1) seg.insert(l, i - 1);
            l = i + 1;
        }
    }
    while (Q--) {
        int c, k;
        cin >> c >> k;
        if (c == 0) seg.insert(k, k);
        if (c == 1) seg.remove(k, k);
        if (c == 2) cout << seg.contains(k) << '\n';
        if (c == 3) cout << seg.lower_bound(k) << '\n';
        if (c == 4) cout << seg.inv_lower_bound(k) << '\n';
    }
}

#line 2 "utilities/integer_segments.hpp"
#include <iostream>
#include <map>
#include <utility>

// CUT begin
// Add/erase ranges on \mathbb{Z}
// Basic implementation idea: https://satanic0258.github.io/snippets/data-structure/SegmentMap.html
template <typename Int> struct integer_segments {
    const Int INVALID = -1;
    Int _sz;
    std::map<Int, Int> mp;
    integer_segments() : _sz(0) {}

    // Get the range [l, r] that satisfies l <= x <= r, or [INVALID, INVALID] otherwise
    std::pair<Int, Int> get(Int x) const {
        auto itr = mp.upper_bound(x);
        if (itr == mp.begin() or (--itr)->second < x) return std::make_pair(INVALID, INVALID);
        return *itr;
    }

    bool contains(Int x) const { return lower_bound(x) == x; }

    // Find the min. y in the set that satisfies x <= y
    Int lower_bound(Int x) const {
        auto itr = mp.upper_bound(x);
        if (itr != mp.begin() and std::prev(itr)->second >= x) return x;
        if (itr != mp.end()) return itr->first;
        return INVALID;
    }

    // Find the max. y in the set that satisfies y <= x
    Int inv_lower_bound(Int x) const {
        auto itr = mp.upper_bound(x);
        if (itr == mp.begin()) return INVALID;
        if ((--itr)->second >= x) return x;
        return itr->second;
    }

    void _mp_upd(Int l, Int r) {
        if (mp.count(l)) _sz -= mp[l] - l + 1;
        mp[l] = r, _sz += r - l + 1;
    }

    // Add {l, l + 1, ..., r} and return the merged segment which [l, r] belongs to
    std::pair<Int, Int> insert(Int l, Int r) {
        auto itl = mp.upper_bound(l), itr = mp.upper_bound(r + 1);
        if (itl != mp.begin() and (--itl)->second < l - 1) { ++itl; }
        if (itl != itr) {
            l = std::min(l, itl->first), r = std::max(r, std::prev(itr)->second);
            for (auto i = itl; i != itr; i++) _sz -= i->second - i->first + 1;
            mp.erase(itl, itr);
        }
        if (mp.count(l)) _sz -= mp[l] - l + 1;
        _mp_upd(l, r);
        return std::make_pair(l, r);
    }

    // Remove {l, l + 1, ..., r}
    void remove(Int l, Int r) {
        auto itl = mp.upper_bound(l), itr = mp.upper_bound(r);
        if (itl != mp.begin() and (--itl)->second < l) { ++itl; }
        if (itl == itr) { return; }
        Int tl = std::min(l, itl->first), tr = std::max(r, std::prev(itr)->second);
        for (auto i = itl; i != itr; i++) _sz -= i->second - i->first + 1;
        mp.erase(itl, itr);
        if (tl < l) _mp_upd(tl, l - 1);
        if (r < tr) _mp_upd(r + 1, tr);
    }

    Int count() const { return _sz; }

    // O((# of segments))
    Int find_by_order(Int k) const {
        if (mp.empty()) return -1;
        Int hi = mp.begin()->first - 1;
        for (const auto &p : mp) {
            const Int l = p.first, r = p.second;
            hi = std::max(hi, l - 1);
            Int add = std::max(Int(0), r - hi);
            if (add >= k + 1) return hi + k + 1;
            k -= add;
            hi = std::max(hi, r);
        }
        return -1;
    }

    // Count elements strictly less than x, O((# of segments))
    Int order_of_key(Int x) const {
        Int ret = 0;
        for (auto p : x) {
            if (p.first >= x) break;
            ret += std::min(x - 1, p.second) - p.first + 1;
        }
        return ret;
    }

    friend std::ostream &operator<<(std::ostream &os, const integer_segments &x) {
        os << '{';
        for (auto &&p : x.mp) os << '[' << p.first << ',' << p.second << "],";
        return os << '}';
    }
};
#line 3 "utilities/test/predecessor_problem.test.cpp"
#include <string>
using namespace std;
#define PROBLEM "https://judge.yosupo.jp/problem/predecessor_problem"

int main() {
    cin.tie(nullptr), ios::sync_with_stdio(false);

    int N, Q;
    string T;
    cin >> N >> Q >> T;
    integer_segments<int> seg;
    T += '0';
    int l = 0;
    for (int i = 0; i < int(T.size()); i++) {
        if (T[i] == '0') {
            if (l <= i - 1) seg.insert(l, i - 1);
            l = i + 1;
        }
    }
    while (Q--) {
        int c, k;
        cin >> c >> k;
        if (c == 0) seg.insert(k, k);
        if (c == 1) seg.remove(k, k);
        if (c == 2) cout << seg.contains(k) << '\n';
        if (c == 3) cout << seg.lower_bound(k) << '\n';
        if (c == 4) cout << seg.inv_lower_bound(k) << '\n';
    }
}

Back to top page