src/merge_sort_tree.hpp

• View this file on GitHub
• Last update: 2025-10-08 00:44:44+09:00
• Include: #include "src/merge_sort_tree.hpp"

Depends on

• src/template.hpp

Verified with

• src/test/aoj_2907.cpp

Code

#pragma once
#include "template.hpp"

template<class S, bool EnableSum = 0> struct merge_sort_tree {
    int n;
    vector<vector<S>> d, s;
    merge_sort_tree(vector<S>& a): n(__bit_ceil(si(a))), d(n * 2), s(n * 2) {
        rep(i, si(a)) d[n + i] = {a[i]}, s[n + i] = {S(), a[i]};
        for(int i = n; --i;) {
            merge(all(d[i * 2]), all(d[i * 2 + 1]), back_inserter(d[i]));
            if(!EnableSum) continue;
            s[i].resize(si(d[i]) + 1, S());
            rep(j, si(d[i])) s[i][j + 1] = s[i][j] + d[i][j];
        }
    }
    int range_cnt(int l, int r, S& u) {
        int ret = 0;
        for(l += n, r += n; l < r; l >>= 1, r >>= 1) {
            if(l & 1) ret += lower_bound(all(d[l]), u) - d[l].begin(), l++;
            if(r & 1) --r, ret += lower_bound(all(d[r]), u) - d[r].begin();
        }
        return ret;
    }
    S range_sum(int l, int r, S& u) {
        assert(EnableSum);
        S ret = 0;
        for(l += n, r += n; l < r; l >>= 1, r >>= 1) {
            if(l & 1) ret += s[l][lower_bound(all(d[l]), u) - d[l].begin()], l++;
            if(r & 1) --r, ret += s[r][lower_bound(all(d[r]), u) - d[r].begin()];
        }
        return ret;
    }
};

#line 2 "src/template.hpp"
#include <bits/stdc++.h>
using namespace std;
#define si(a) (long)a.size()
#define fi first
#define se second
#define all(x) x.begin(), x.end()
#define rep(i, n) for(long i = 0; i < (long)(n); ++i)
template<typename T> bool chmin(T& a, T b) { return b < a ? (a = b, 1) : 0; }
template<typename T> bool chmax(T& a, T b) { return b > a ? (a = b, 1) : 0; }
struct _ {
    _() { cin.tie(0)->sync_with_stdio(0), cout.tie(0), cout << fixed << setprecision(16); }
} __;
#line 3 "src/merge_sort_tree.hpp"

template<class S, bool EnableSum = 0> struct merge_sort_tree {
    int n;
    vector<vector<S>> d, s;
    merge_sort_tree(vector<S>& a): n(__bit_ceil(si(a))), d(n * 2), s(n * 2) {
        rep(i, si(a)) d[n + i] = {a[i]}, s[n + i] = {S(), a[i]};
        for(int i = n; --i;) {
            merge(all(d[i * 2]), all(d[i * 2 + 1]), back_inserter(d[i]));
            if(!EnableSum) continue;
            s[i].resize(si(d[i]) + 1, S());
            rep(j, si(d[i])) s[i][j + 1] = s[i][j] + d[i][j];
        }
    }
    int range_cnt(int l, int r, S& u) {
        int ret = 0;
        for(l += n, r += n; l < r; l >>= 1, r >>= 1) {
            if(l & 1) ret += lower_bound(all(d[l]), u) - d[l].begin(), l++;
            if(r & 1) --r, ret += lower_bound(all(d[r]), u) - d[r].begin();
        }
        return ret;
    }
    S range_sum(int l, int r, S& u) {
        assert(EnableSum);
        S ret = 0;
        for(l += n, r += n; l < r; l >>= 1, r >>= 1) {
            if(l & 1) ret += s[l][lower_bound(all(d[l]), u) - d[l].begin()], l++;
            if(r & 1) --r, ret += s[r][lower_bound(all(d[r]), u) - d[r].begin()];
        }
        return ret;
    }
};

Back to top page