shogo-library

競技プログラミング用の C++ ライブラリ

View the Project on GitHub shogo314/shogo-library

:heavy_check_mark: test/unittest/counter-segtree.test.cpp

Depends on

Code

#define PROBLEM "https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ITP1_1_A"

// CounterSegTree を愚直な実装とランダムに比較する
// 問題の入力は使わず、すべて一致したら Hello World を出力する

#include <iostream>
#include <random>
#include <utility>
#include <vector>

#include "../../src/data_structures/counter-segtree.hpp"

int main() {
    std::mt19937 rng(314);
    bool ok = true;
    for (int rep = 0; rep < 2000 && ok; rep++) {
        int n = rng() % 9 + 1;
        std::vector<int> v(n);
        for (int& x : v) x = rng() % 6;
        CounterSegTree<int> seg(v);
        // key_type / mapped_type は外から使える
        static_assert(std::is_same_v<CounterSegTree<int>::key_type, int>);
        static_assert(std::is_same_v<CounterSegTree<int>::mapped_type, int>);
        for (int t = 0; t < 20; t++) {
            if (rng() % 2) {
                int k = rng() % n, x = rng() % 6;
                seg.set(k, x);
                v[k] = x;
            }
            for (int k = 0; k < n; k++) {
                if (seg.get(k) != v[k]) ok = false;
            }
            int l = rng() % (n + 1), r = rng() % (n + 1);
            if (l > r) std::swap(l, r);
            int x = rng() % 7;
            // lower_bound: [l, r) にある x 以上の最小の値とその個数 (なければ個数 0)
            int best = -1, cnt = 0, same = 0;
            for (int i = l; i < r; i++) {
                if (v[i] >= x && (best == -1 || v[i] < best)) best = v[i];
                if (v[i] == x) same++;
            }
            for (int i = l; i < r; i++) {
                if (v[i] == best) cnt++;
            }
            auto lb = seg.lower_bound(l, r, x);
            if (lb.second != cnt) ok = false;
            if (cnt > 0 && lb.first != best) ok = false;
            // find: [l, r) にある x の個数
            auto f = seg.find(l, r, x);
            if (f.first != x || f.second != same) ok = false;
        }
    }
    if (ok) std::cout << "Hello World" << std::endl;
    else std::cerr << "failed" << std::endl;
}
#line 1 "test/unittest/counter-segtree.test.cpp"
#define PROBLEM "https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ITP1_1_A"

// CounterSegTree を愚直な実装とランダムに比較する
// 問題の入力は使わず、すべて一致したら Hello World を出力する

#include <iostream>
#include <random>
#include <utility>
#include <vector>

#line 2 "src/data_structures/counter-segtree.hpp"

/**
 * @file counter-segtree.hpp
 * @brief 区間内の値の個数を数えるセグメント木
 */

#include <cassert>
#include <map>
#line 12 "src/data_structures/counter-segtree.hpp"

template <typename T>
class CounterSegTree {
   public:
    using key_type = T;
    using mapped_type = int;

   private:
    int n, sz, height;
    std::vector<std::map<key_type, mapped_type>> data;

   public:
    explicit CounterSegTree(const std::vector<key_type>& v) {
        n = v.size();
        sz = 1;
        height = 0;
        while (sz < n) {
            sz <<= 1;
            height++;
        }
        data.assign(sz * 2, {});
        for (int i = 0; i < n; i++) {
            data[i + sz][v[i]] = 1;
        }
        for (int i = sz - 1; i > 0; i--) {
            data[i] = data[2 * i];
            for (auto [s, c] : data[2 * i + 1]) {
                data[i][s] += c;
            }
        }
    }
    key_type get(int k) const {
        assert(0 <= k and k < n);
        return data[k + sz].begin()->first;
    }
    void set(int k, const key_type& x) {
        assert(0 <= k and k < n);
        k += sz;
        key_type y = data[k].begin()->first;
        while (k > 0) {
            auto itr = data[k].find(y);
            if (itr->second == 1) {
                data[k].erase(itr);
            } else {
                itr->second--;
            }
            data[k][x]++;
            k >>= 1;
        }
    }
    std::pair<key_type, mapped_type> lower_bound(int l, int r, key_type x) const {
        std::pair<key_type, mapped_type> ret = {{}, 0};
        l += sz;
        r += sz;
        while (l < r) {
            if (l & 1) {
                auto itr = data[l].lower_bound(x);
                if (itr != data[l].end()) {
                    if (ret.second == 0) {
                        ret = *itr;
                    } else if (itr->first == ret.first) {
                        ret.second += itr->second;
                    } else if (itr->first < ret.first) {
                        ret = *itr;
                    }
                }
                ++l;
            }
            if (r & 1) {
                --r;
                auto itr = data[r].lower_bound(x);
                if (itr != data[r].end()) {
                    if (ret.second == 0) {
                        ret = *itr;
                    } else if (itr->first == ret.first) {
                        ret.second += itr->second;
                    } else if (itr->first < ret.first) {
                        ret = *itr;
                    }
                }
            }
            l >>= 1;
            r >>= 1;
        }
        return ret;
    }
    std::pair<key_type, mapped_type> find(int l, int r, key_type x) const {
        std::pair<key_type, mapped_type> ret = {x, 0};
        l += sz;
        r += sz;
        while (l < r) {
            if (l & 1) {
                auto itr = data[l].find(x);
                if (itr != data[l].end()) {
                    ret.second += itr->second;
                }
                ++l;
            }
            if (r & 1) {
                --r;
                auto itr = data[r].find(x);
                if (itr != data[r].end()) {
                    ret.second += itr->second;
                }
            }
            l >>= 1;
            r >>= 1;
        }
        return ret;
    }
};
#line 12 "test/unittest/counter-segtree.test.cpp"

int main() {
    std::mt19937 rng(314);
    bool ok = true;
    for (int rep = 0; rep < 2000 && ok; rep++) {
        int n = rng() % 9 + 1;
        std::vector<int> v(n);
        for (int& x : v) x = rng() % 6;
        CounterSegTree<int> seg(v);
        // key_type / mapped_type は外から使える
        static_assert(std::is_same_v<CounterSegTree<int>::key_type, int>);
        static_assert(std::is_same_v<CounterSegTree<int>::mapped_type, int>);
        for (int t = 0; t < 20; t++) {
            if (rng() % 2) {
                int k = rng() % n, x = rng() % 6;
                seg.set(k, x);
                v[k] = x;
            }
            for (int k = 0; k < n; k++) {
                if (seg.get(k) != v[k]) ok = false;
            }
            int l = rng() % (n + 1), r = rng() % (n + 1);
            if (l > r) std::swap(l, r);
            int x = rng() % 7;
            // lower_bound: [l, r) にある x 以上の最小の値とその個数 (なければ個数 0)
            int best = -1, cnt = 0, same = 0;
            for (int i = l; i < r; i++) {
                if (v[i] >= x && (best == -1 || v[i] < best)) best = v[i];
                if (v[i] == x) same++;
            }
            for (int i = l; i < r; i++) {
                if (v[i] == best) cnt++;
            }
            auto lb = seg.lower_bound(l, r, x);
            if (lb.second != cnt) ok = false;
            if (cnt > 0 && lb.first != best) ok = false;
            // find: [l, r) にある x の個数
            auto f = seg.find(l, r, x);
            if (f.first != x || f.second != same) ok = false;
        }
    }
    if (ok) std::cout << "Hello World" << std::endl;
    else std::cerr << "failed" << std::endl;
}
Back to top page