重み付きウェーブレット行列 (lib/data_structure/wavelet_matrix_rectangle_sum.hpp)

• View this file on GitHub
• Last update: 2025-06-10 12:07:12+09:00
• Include: #include "lib/data_structure/wavelet_matrix_rectangle_sum.hpp"

Depends on

• 完備辞書 (lib/internal/internal_bit_vector.hpp)

Verified with

• test/aoj/challenge/2270.test.cpp
• test/yosupo/data_structure/rectangle_sum.test.cpp

Code

#pragma once
#include <cassert>
#include <numeric>
#include <tuple>
#include <vector>
#include "internal/internal_bit_vector.hpp"

/// @brief 重み付きウェーブレット行列
template <class T, class U = T, int L = 30>
struct wavelet_matrix_rectangle_sum {
    wavelet_matrix_rectangle_sum() = default;
    template <class Value>
    wavelet_matrix_rectangle_sum(const std::vector<T> &v, const std::vector<Value> &u) : length(v.size()) {
        assert(v.size() == u.size());
        std::vector<int> l(length), r(length), ord(length);
        std::iota(ord.begin(), ord.end(), 0);
        for (int level = L - 1; level >= 0; level--) {
            matrix[level] = internal::bit_vector(length + 1);
            int left = 0, right = 0;
            for (int i = 0; i < length; i++) {
                if ((v[ord[i]] >> level) & 1) {
                    matrix[level].set(i);
                    r[right++] = ord[i];
                } else {
                    l[left++] = ord[i];
                }
            }
            mid[level] = left;
            matrix[level].build();
            ord.swap(l);
            for (int i = 0; i < right; i++) ord[left + i] = r[i];
            cs[level].resize(length + 1);
            cs[level][0] = U(0);
            for (int i = 0; i < length; i++) cs[level][i + 1] = cs[level][i] + u[ord[i]];
        }
    }

    U kth_smallest_sum(int l, int r, int k) const {
        assert(0 <= k && k <= r - l);
        T val = T();
        U res = U();
        for (int level = L - 1; level >= 0; --level) {
            int cnt = matrix[level].rank(false, r) - matrix[level].rank(false, l);
            bool f = cnt <= k;
            if (f) {
                val |= T(1) << level;
                res += cs[level][matrix[level].rank(false, r)] - cs[level][matrix[level].rank(false, l)];
                k -= cnt;
            }
            std::tie(l, r) = succ(f, l, r, level);
        }
        return res + val * k;
    }

    U kth_largest_sum(int l, int r, int k) const {
        return cs[L - 1][matrix[L - 1].rank(false, r)] + cs[L - 1][matrix[L - 1].rank(true, r)] -
               cs[L - 1][matrix[L - 1].rank(false, l)] - cs[L - 1][matrix[L - 1].rank(true, l)] -
               kth_smallest_sum(l, r, r - l - k);
    }

    U range_sum(int r, T x) const { return range_sum(0, r, x); }

    U range_sum(int l, int r, T x) const {
        for (int level = L - 1; level >= 0; level--) std::tie(l, r) = succ((x >> level) & 1, l, r, level);
        return cs[0][matrix[0].rank(false, r)] - cs[0][matrix[0].rank(false, l)];
    }

    U rect_sum(int l, int r, T upper) const {
        U res = 0;
        for (int level = L - 1; level >= 0; level--) {
            bool f = (upper >> level) & 1;
            if (f) res += cs[level][matrix[level].rank(false, r)] - cs[level][matrix[level].rank(false, l)];
            std::tie(l, r) = succ(f, l, r, level);
        }
        return res;
    }

    U rect_sum(int l, int r, T lower, T upper) const { return rect_sum(l, r, upper) - rect_sum(l, r, lower); }

  private:
    int length;
    internal::bit_vector matrix[L];
    int mid[L];
    std::vector<U> cs[L];

    std::pair<int, int> succ(bool f, int l, int r, int level) const {
        return {matrix[level].rank(f, l) + mid[level] * f, matrix[level].rank(f, r) + mid[level] * f};
    }
};

Traceback (most recent call last):
  File "/home/runner/.local/lib/python3.12/site-packages/competitive_verifier/oj/resolver.py", line 291, in resolve
    bundled_code = language.bundle(path, basedir=basedir)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/runner/.local/lib/python3.12/site-packages/competitive_verifier/oj/verify/languages/cplusplus.py", line 242, in bundle
    bundler.update(path)
  File "/home/runner/.local/lib/python3.12/site-packages/competitive_verifier/oj/verify/languages/cplusplus_bundle.py", line 479, in update
    self._resolve(pathlib.Path(included), included_from=path)
  File "/home/runner/.local/lib/python3.12/site-packages/competitive_verifier/oj/verify/languages/cplusplus_bundle.py", line 286, in _resolve
    raise BundleErrorAt(path, -1, "no such header")
competitive_verifier.oj.verify.languages.cplusplus_bundle.BundleErrorAt: internal/internal_bit_vector.hpp: line -1: no such header

Back to top page