shogo-library

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

View the Project on GitHub shogo314/shogo-library

:heavy_check_mark: マージソートツリー
(src/data_structures/merge-sort-tree.hpp)

各要素が値 value[i] とキー key[i] を持つ列に対して、「添字が $[l, r)$ で、キーが $[a, b)$ に入る要素の値の積」を求めます。更新はできません。

std::vector<long long> value = {1, 2, 3, 4, 5};
std::vector<long long> key   = {5, 1, 4, 2, 3};
MergeSortTree<> mst(value, key);           // 値は long long の和
mst.prod(1, 4, 2, 5);                      // i ∈ [1, 4) かつ key ∈ [2, 5) → 3 + 4 = 7
mst.prod(0, 5, std::nullopt, 3);           // key < 3 → 2 + 4 = 6

点の重みの矩形和 (x を添字、y をキーにする) や、区間内で閾値未満の値の個数・和などに使えます。

テンプレート引数

MergeSortTree<G = group::Sum<long long>, K = long long, Comp = std::less<K>>

引数 説明
G 値の可換群 (algebra/group.hpp)。逆元を使って積を求める
K キーの型
Comp キーの比較関数

メンバ関数

関数 説明 計算量
MergeSortTree(value, key, g = G()) 値の列とキーの列から作る $O(n \log n)$
MergeSortTree(value_key, g = G()) (値, キー) の pair の列から作る $O(n \log n)$
prod(l, r, a, b) $l \le i < r$ かつ $a \le \mathrm{key}_i < b$ の $\mathrm{value}_i$ の積。どの引数も std::nullopt (省略) にでき、l、r は 0、n、a、b は $-\infty$、$+\infty$ になる $O(\log^2 n)$

関連

Depends on

Verified with

Code

#pragma once

/**
 * @file merge-sort-tree.hpp
 * @brief マージソートツリー
 */
#include <algorithm>
#include <cassert>
#include <functional>
#include <optional>
#include <utility>
#include <vector>
#include "../algebra/group.hpp"

/**
 * @brief 区間の閾値以内の値の積
 *
 * @tparam G 値の可換群 (algebra/group.hpp)
 * @tparam K key_type ソートに使う型
 * @tparam Comp Kを比較するファンクタ
 */
template <typename G = group::Sum<long long>,
          typename K = long long,
          class Comp = std::less<K>>
class MergeSortTree {
   public:
    using value_type = typename G::S;
    using key_type = K;

    inline constexpr static auto comp = Comp();

   private:
    G g;
    value_type op(const value_type& x, const value_type& y) const { return g.op(x, y); }
    value_type e() const { return g.e(); }
    value_type inv(const value_type& x) const { return g.inv(x); }

   private:
    int n, sz, height;
    std::vector<key_type> key_data;
    std::vector<value_type> cumulative_value;

    void initialize(const std::vector<value_type>& value, const std::vector<key_type>& key) {
        n = key.size();
        sz = 1;
        height = 1;
        while (sz < n) {
            sz <<= 1;
            height++;
        }
        key_data.assign(sz * height, {});
        std::vector<value_type> value_data(sz * height, e());
        cumulative_value.assign(sz * height, {});
        for (int i = 0; i < n; i++) {
            key_data[(height - 1) * sz + i] = key[i];
            value_data[(height - 1) * sz + i] = value[i];
            cumulative_value[(height - 1) * sz + i] = value[i];
        }
        int t = 1;
        for (int h = height - 1; h > 0; h--) {
            for (int i = 0; i < n; i += t * 2) {
                int j1 = h * sz + i;
                int j2 = h * sz + std::min(n, i + t);
                int j0 = (h - 1) * sz + i;
                int last1 = j2;
                int last2 = h * sz + std::min(n, i + t * 2);
                while (j1 != last1 || j2 != last2) {
                    if (j1 == last1 || (j2 < last2 && comp(key_data[j2], key_data[j1]))) {
                        key_data[j0] = key_data[j2];
                        value_data[j0] = value_data[j2];
                        j0++;
                        j2++;
                    } else {
                        key_data[j0] = key_data[j1];
                        value_data[j0] = value_data[j1];
                        j0++;
                        j1++;
                    }
                }
                cumulative_value[(h - 1) * sz + i] = value_data[(h - 1) * sz + i];
                for (int j = i + 1; j < std::min(n, i + t * 2); j++) {
                    cumulative_value[(h - 1) * sz + j] = op(cumulative_value[(h - 1) * sz + j - 1], value_data[(h - 1) * sz + j]);
                }
            }
            t <<= 1;
        }
    }

    value_type prod_section(int l, int r, std::optional<key_type> a, std::optional<key_type> b) const {
        value_type ret = cumulative_value[r - 1];
        if (b.has_value()) {
            int i = std::lower_bound(key_data.begin() + l, key_data.begin() + r, b.value(), comp) - key_data.begin();
            if (i != l) {
                ret = cumulative_value[i - 1];
            } else {
                ret = e();
            }
        }
        if (a.has_value()) {
            int i = std::lower_bound(key_data.begin() + l, key_data.begin() + r, a.value(), comp) - key_data.begin();
            if (i != l) {
                ret = op(ret, inv(cumulative_value[i - 1]));
            }
        }
        return ret;
    }

   public:
    MergeSortTree() = default;
    /**
     * @param value_key valueとkeyのpairのvector
     */
    explicit MergeSortTree(const std::vector<std::pair<value_type, key_type>>& value_key, G g = G()) : g(std::move(g)) {
        std::vector<key_type> key;
        std::vector<value_type> value;
        key.reserve(value_key.size());
        value.reserve(value_key.size());
        for (size_t i = 0; i < value_key.size(); i++) {
            value.push_back(value_key[i].first);
            key.push_back(value_key[i].second);
        }
        this->initialize(value, key);
    }
    /**
     * @param value prodで計算する対象
     * @param key ソートする基準
     */
    MergeSortTree(const std::vector<value_type>& value, const std::vector<key_type>& key, G g = G()) : g(std::move(g)) {
        assert(key.size() == value.size());
        this->initialize(value, key);
    }

    /**
     * @brief product value[i] s.t. a <= key[i] < b , i in [l, r)
     *
     * @param l 半開区間の開始 nulloptの場合は0
     * @param r 半開区間の終端 0<=l<=r<=n nulloptの場合はn
     * @param a nulloptの場合は負の無限大
     * @param b nulloptの場合は正の無限大
     */
    value_type prod(std::optional<int> l_ = std::nullopt, std::optional<int> r_ = std::nullopt, std::optional<key_type> a = std::nullopt, std::optional<key_type> b = std::nullopt) const {
        int l = l_.value_or(0), r = r_.value_or(n);
        assert(0 <= l && l <= r && r <= n);
        if (a.has_value() && b.has_value() && !comp(a.value(), b.value())) return e();
        value_type ret = e();
        int h = height - 1;
        int t = 1;
        while (l < r) {
            if (l & t) {
                ret = op(ret, prod_section(h * sz + l, h * sz + l + t, a, b));
                l += t;
            }
            if (r & t) {
                r -= t;
                ret = op(ret, prod_section(h * sz + r, h * sz + r + t, a, b));
            }
            h--;
            t <<= 1;
        }
        return ret;
    }
};
#line 2 "src/data_structures/merge-sort-tree.hpp"

/**
 * @file merge-sort-tree.hpp
 * @brief マージソートツリー
 */
#include <algorithm>
#include <cassert>
#include <functional>
#include <optional>
#include <utility>
#include <vector>
#line 2 "src/algebra/group.hpp"

/**
 * @file group.hpp
 * @brief 群の定義
 *
 * 群はモノイド (algebra/monoid.hpp) のメンバに加えて次のメンバを持つ型として表す
 * - S inv(const S& a): 逆元
 */

#line 2 "src/algebra/monoid.hpp"

/**
 * @file monoid.hpp
 * @brief モノイドの定義
 *
 * モノイドは次のメンバを持つ型として表す
 * - using S: 要素の型
 * - S op(const S& a, const S& b): 積 (結合的)
 * - S e(): 単位元
 * op, e は static でも非 static でもよい (データ構造はモノイドのオブジェクトを保持して呼び出す)
 */

#line 15 "src/algebra/monoid.hpp"
#include <istream>
#include <limits>
#include <numeric>
#include <ostream>
#include <type_traits>
#line 21 "src/algebra/monoid.hpp"

/**
 * @brief 1次関数 f(x) = a * x + b
 *
 * monoid::Composite の要素や acted_monoid::AffineSum の作用素に使う
 * 例: Affine<mint> f{2, 3}; f(x) は 2x + 3。省略すると恒等写像 {1, 0}
 */
template <typename T>
struct Affine {
    T a = T(1);  //!< 傾き
    T b = T(0);  //!< 切片
    constexpr T operator()(const T& x) const { return a * x + b; }
    friend constexpr bool operator==(const Affine& f, const Affine& g) { return f.a == g.a && f.b == g.b; }
    friend constexpr bool operator!=(const Affine& f, const Affine& g) { return !(f == g); }
    template <class CharT, class Traits>
    friend std::basic_istream<CharT, Traits>& operator>>(std::basic_istream<CharT, Traits>& is, Affine& f) {
        return is >> f.a >> f.b;
    }
    template <class CharT, class Traits>
    friend std::basic_ostream<CharT, Traits>& operator<<(std::basic_ostream<CharT, Traits>& os, const Affine& f) {
        return os << f.a << CharT(' ') << f.b;
    }
};

namespace monoid {

/**
 * @brief 和
 */
template <typename T>
struct Sum {
    using S = T;
    static constexpr S op(const S& a, const S& b) { return a + b; }
    static constexpr S e() { return S(0); }
};

/**
 * @brief 積
 */
template <typename T>
struct Prod {
    using S = T;
    static constexpr S op(const S& a, const S& b) { return a * b; }
    static constexpr S e() { return S(1); }
};

/**
 * @brief 最大値 単位元は型の最小値
 */
template <typename T>
struct Max {
    using S = T;
    static constexpr S op(const S& a, const S& b) { return std::max(a, b); }
    static constexpr S e() { return std::numeric_limits<S>::lowest(); }
};

/**
 * @brief 最小値 単位元は型の最大値
 */
template <typename T>
struct Min {
    using S = T;
    static constexpr S op(const S& a, const S& b) { return std::min(a, b); }
    static constexpr S e() { return std::numeric_limits<S>::max(); }
};

/**
 * @brief 最小値と最大値の組 {min, max}
 */
template <typename T>
struct MinMax {
    using S = std::pair<T, T>;
    static constexpr S op(const S& a, const S& b) { return {std::min(a.first, b.first), std::max(a.second, b.second)}; }
    static constexpr S e() { return {std::numeric_limits<T>::max(), std::numeric_limits<T>::lowest()}; }
};

/**
 * @brief 排他的論理和
 */
template <typename T>
struct Xor {
    using S = T;
    static constexpr S op(const S& a, const S& b) { return a ^ b; }
    static constexpr S e() { return S(0); }
};

/**
 * @brief 論理積
 */
template <typename T>
struct And {
    using S = T;
    static constexpr S op(const S& a, const S& b) { return a & b; }
    static constexpr S e() { return ~S(0); }
};

/**
 * @brief 論理和
 */
template <typename T>
struct Or {
    using S = T;
    static constexpr S op(const S& a, const S& b) { return a | b; }
    static constexpr S e() { return S(0); }
};

/**
 * @brief 最大公約数 単位元は0
 */
template <typename T>
struct Gcd {
    static_assert(std::is_integral_v<T>);
    using S = T;
    static constexpr S op(const S& a, const S& b) { return std::gcd(a, b); }
    static constexpr S e() { return S(0); }
};

/**
 * @brief 最小公倍数 単位元は1
 */
template <typename T>
struct Lcm {
    static_assert(std::is_integral_v<T>);
    using S = T;
    static constexpr S op(const S& a, const S& b) { return std::lcm(a, b); }
    static constexpr S e() { return S(1); }
};

/**
 * @brief 1次関数 (Affine) の合成
 * op(f, g) = g∘f (左から順に適用する)
 * 区間 [l, r) の積は f_{r-1}(...f_l(x)) を表す
 */
template <typename T>
struct Composite {
    using S = Affine<T>;
    static constexpr S op(const S& f, const S& g) { return {g.a * f.a, g.a * f.b + g.b}; }
    static constexpr S e() { return {T(1), T(0)}; }
};

/**
 * @brief 括弧列 対応の取れていない ')' と '(' の個数
 */
struct Parenthesis {
    struct S {
        long long close; // 対応の取れていない ')' の個数
        long long open;  // 対応の取れていない '(' の個数
        constexpr S(char c = 0) : close(c == ')'), open(c == '(') {}
        constexpr S(long long close, long long open) : close(close), open(open) {}
        /**
         * @brief 正しい括弧列か
         */
        constexpr bool is_valid() const { return close == 0 && open == 0; }
    };
    static constexpr S op(const S& a, const S& b) {
        if(a.open >= b.close) return S(a.close, a.open - b.close + b.open);
        return S(a.close + b.close - a.open, b.open);
    }
    static constexpr S e() { return S(); }
};

/**
 * @brief 最小値とその位置 同じ値なら左側を優先
 *
 * @tparam T 値の型
 * @tparam U 位置などの付随する値の型
 */
template <typename T, typename U = int>
struct MinWith {
    using S = std::pair<T, U>;
    static constexpr S op(const S& a, const S& b) { return b.first < a.first ? b : a; }
    static constexpr S e() { return {std::numeric_limits<T>::max(), U{}}; }
};

/**
 * @brief 最大値とその位置 同じ値なら左側を優先
 *
 * @tparam T 値の型
 * @tparam U 位置などの付随する値の型
 */
template <typename T, typename U = int>
struct MaxWith {
    using S = std::pair<T, U>;
    static constexpr S op(const S& a, const S& b) { return a.first < b.first ? b : a; }
    static constexpr S e() { return {std::numeric_limits<T>::lowest(), U{}}; }
};

/**
 * @brief 最小値とその個数
 *
 * @tparam T 値の型
 * @tparam U 個数の型
 */
template <typename T, typename U = long long>
struct MinCount {
    using S = std::pair<T, U>;
    static constexpr S op(const S& a, const S& b) {
        if(a.first < b.first) return a;
        if(b.first < a.first) return b;
        return {a.first, a.second + b.second};
    }
    static constexpr S e() { return {std::numeric_limits<T>::max(), U(0)}; }
};

/**
 * @brief 最大値とその個数
 *
 * @tparam T 値の型
 * @tparam U 個数の型
 */
template <typename T, typename U = long long>
struct MaxCount {
    using S = std::pair<T, U>;
    static constexpr S op(const S& a, const S& b) {
        if(a.first < b.first) return b;
        if(b.first < a.first) return a;
        return {a.first, a.second + b.second};
    }
    static constexpr S e() { return {std::numeric_limits<T>::lowest(), U(0)}; }
};

} // namespace monoid

/**
 * @brief 関数オブジェクトと単位元から実行時に作るモノイド
 * ラムダ式を使える
 *
 * 例: SegTree seg(v, Monoid([](int a, int b) { return std::max(a, b); }, -1));
 *
 * @tparam T 要素の型
 * @tparam Op 積の関数オブジェクトの型
 */
template <typename T, typename Op>
struct Monoid {
    using S = T;
    Op op_object;
    S identity;
    Monoid() = default;
    constexpr Monoid(Op op, const S& e) : op_object(std::move(op)), identity(e) {}
    constexpr S op(const S& a, const S& b) const { return op_object(a, b); }
    constexpr S e() const { return identity; }
};
template <typename Op, typename T>
Monoid(Op, T) -> Monoid<T, Op>;
#line 12 "src/algebra/group.hpp"

namespace group {

/**
 * @brief 加法群
 */
template <typename T>
struct Sum : monoid::Sum<T> {
    using S = T;
    static constexpr S inv(const S& a) { return -a; }
};

/**
 * @brief 乗法群 (0を含まない体の元など)
 */
template <typename T>
struct Prod : monoid::Prod<T> {
    using S = T;
    static constexpr S inv(const S& a) { return S(1) / a; }
};

/**
 * @brief 排他的論理和
 */
template <typename T>
struct Xor : monoid::Xor<T> {
    using S = T;
    static constexpr S inv(const S& a) { return a; }
};

} // namespace group

/**
 * @brief 関数オブジェクトと単位元から実行時に作る群
 * ラムダ式を使える
 *
 * @tparam T 要素の型
 * @tparam Op 積の関数オブジェクトの型
 * @tparam Inv 逆元の関数オブジェクトの型
 */
template <typename T, typename Op, typename Inv>
struct Group {
    using S = T;
    Op op_object;
    S identity;
    Inv inv_object;
    Group() = default;
    constexpr Group(Op op, const S& e, Inv inv) : op_object(std::move(op)), identity(e), inv_object(std::move(inv)) {}
    constexpr S op(const S& a, const S& b) const { return op_object(a, b); }
    constexpr S e() const { return identity; }
    constexpr S inv(const S& a) const { return inv_object(a); }
};
template <typename Op, typename T, typename Inv>
Group(Op, T, Inv) -> Group<T, Op, Inv>;
#line 14 "src/data_structures/merge-sort-tree.hpp"

/**
 * @brief 区間の閾値以内の値の積
 *
 * @tparam G 値の可換群 (algebra/group.hpp)
 * @tparam K key_type ソートに使う型
 * @tparam Comp Kを比較するファンクタ
 */
template <typename G = group::Sum<long long>,
          typename K = long long,
          class Comp = std::less<K>>
class MergeSortTree {
   public:
    using value_type = typename G::S;
    using key_type = K;

    inline constexpr static auto comp = Comp();

   private:
    G g;
    value_type op(const value_type& x, const value_type& y) const { return g.op(x, y); }
    value_type e() const { return g.e(); }
    value_type inv(const value_type& x) const { return g.inv(x); }

   private:
    int n, sz, height;
    std::vector<key_type> key_data;
    std::vector<value_type> cumulative_value;

    void initialize(const std::vector<value_type>& value, const std::vector<key_type>& key) {
        n = key.size();
        sz = 1;
        height = 1;
        while (sz < n) {
            sz <<= 1;
            height++;
        }
        key_data.assign(sz * height, {});
        std::vector<value_type> value_data(sz * height, e());
        cumulative_value.assign(sz * height, {});
        for (int i = 0; i < n; i++) {
            key_data[(height - 1) * sz + i] = key[i];
            value_data[(height - 1) * sz + i] = value[i];
            cumulative_value[(height - 1) * sz + i] = value[i];
        }
        int t = 1;
        for (int h = height - 1; h > 0; h--) {
            for (int i = 0; i < n; i += t * 2) {
                int j1 = h * sz + i;
                int j2 = h * sz + std::min(n, i + t);
                int j0 = (h - 1) * sz + i;
                int last1 = j2;
                int last2 = h * sz + std::min(n, i + t * 2);
                while (j1 != last1 || j2 != last2) {
                    if (j1 == last1 || (j2 < last2 && comp(key_data[j2], key_data[j1]))) {
                        key_data[j0] = key_data[j2];
                        value_data[j0] = value_data[j2];
                        j0++;
                        j2++;
                    } else {
                        key_data[j0] = key_data[j1];
                        value_data[j0] = value_data[j1];
                        j0++;
                        j1++;
                    }
                }
                cumulative_value[(h - 1) * sz + i] = value_data[(h - 1) * sz + i];
                for (int j = i + 1; j < std::min(n, i + t * 2); j++) {
                    cumulative_value[(h - 1) * sz + j] = op(cumulative_value[(h - 1) * sz + j - 1], value_data[(h - 1) * sz + j]);
                }
            }
            t <<= 1;
        }
    }

    value_type prod_section(int l, int r, std::optional<key_type> a, std::optional<key_type> b) const {
        value_type ret = cumulative_value[r - 1];
        if (b.has_value()) {
            int i = std::lower_bound(key_data.begin() + l, key_data.begin() + r, b.value(), comp) - key_data.begin();
            if (i != l) {
                ret = cumulative_value[i - 1];
            } else {
                ret = e();
            }
        }
        if (a.has_value()) {
            int i = std::lower_bound(key_data.begin() + l, key_data.begin() + r, a.value(), comp) - key_data.begin();
            if (i != l) {
                ret = op(ret, inv(cumulative_value[i - 1]));
            }
        }
        return ret;
    }

   public:
    MergeSortTree() = default;
    /**
     * @param value_key valueとkeyのpairのvector
     */
    explicit MergeSortTree(const std::vector<std::pair<value_type, key_type>>& value_key, G g = G()) : g(std::move(g)) {
        std::vector<key_type> key;
        std::vector<value_type> value;
        key.reserve(value_key.size());
        value.reserve(value_key.size());
        for (size_t i = 0; i < value_key.size(); i++) {
            value.push_back(value_key[i].first);
            key.push_back(value_key[i].second);
        }
        this->initialize(value, key);
    }
    /**
     * @param value prodで計算する対象
     * @param key ソートする基準
     */
    MergeSortTree(const std::vector<value_type>& value, const std::vector<key_type>& key, G g = G()) : g(std::move(g)) {
        assert(key.size() == value.size());
        this->initialize(value, key);
    }

    /**
     * @brief product value[i] s.t. a <= key[i] < b , i in [l, r)
     *
     * @param l 半開区間の開始 nulloptの場合は0
     * @param r 半開区間の終端 0<=l<=r<=n nulloptの場合はn
     * @param a nulloptの場合は負の無限大
     * @param b nulloptの場合は正の無限大
     */
    value_type prod(std::optional<int> l_ = std::nullopt, std::optional<int> r_ = std::nullopt, std::optional<key_type> a = std::nullopt, std::optional<key_type> b = std::nullopt) const {
        int l = l_.value_or(0), r = r_.value_or(n);
        assert(0 <= l && l <= r && r <= n);
        if (a.has_value() && b.has_value() && !comp(a.value(), b.value())) return e();
        value_type ret = e();
        int h = height - 1;
        int t = 1;
        while (l < r) {
            if (l & t) {
                ret = op(ret, prod_section(h * sz + l, h * sz + l + t, a, b));
                l += t;
            }
            if (r & t) {
                r -= t;
                ret = op(ret, prod_section(h * sz + r, h * sz + r + t, a, b));
            }
            h--;
            t <<= 1;
        }
        return ret;
    }
};
Back to top page