shogo-library

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

View the Project on GitHub shogo314/shogo-library

:heavy_check_mark: 区間内の値の個数を数えるセグメント木
(src/data_structures/counter-segtree.hpp)

各ノードに部分列の値の個数 (std::map) を持つセグメント木です。1点更新と、区間内の特定の値の個数・x 以上の最小の値を求められます。

CounterSegTree<int> cs({3, 1, 3, 2, 3});
cs.find(0, 4, 3);         // {3, 2}: [0, 4) に 3 は 2 個
cs.lower_bound(1, 5, 2);  // {2, 1}: [1, 5) で 2 以上の最小の値は 2 で、1 個
cs.set(2, 5);

メンバ関数

関数 説明 計算量
CounterSegTree(v) 列 v から作る $O(n \log^2 n)$
get(k) k 番目の値 $O(1)$
set(k, x) k 番目を x に更新 $O(\log^2 n)$
find(l, r, x) {x, 区間 [l, r) の x の個数} $O(\log^2 n)$
lower_bound(l, r, x) 区間 $[l, r)$ で x 以上の最小の値と、その個数の pair。なければ個数が 0 $O(\log^2 n)$

メモリは $O(n \log n)$ 使います。

Verified with

Code

#pragma once

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

#include <cassert>
#include <map>
#include <utility>
#include <vector>

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 2 "src/data_structures/counter-segtree.hpp"

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

#include <cassert>
#include <map>
#include <utility>
#include <vector>

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;
    }
};
Back to top page