library

This documentation is automatically generated by online-judge-tools/verification-helper

View the Project on GitHub gyouzasushi/library

:heavy_check_mark: test/unit/wavelet_matrix.test.cpp

Depends on

Code

#define PROBLEM "http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ITP1_1_A"
#include <algorithm>
#include <cassert>
#include <iostream>
#include <random>
#include <vector>

#include "datastructure/wavelet_matrix.hpp"

template <class T, class S>
void check(const std::vector<T>& a, const std::vector<S>& w, std::mt19937& rng, T lo_v, T hi_v) {
    int n = int(a.size());
    wavelet_matrix<T, S> wm(a, w);
    wavelet_matrix<T> wc(a);
    assert(wm.size() == n && wc.size() == n);
    auto rnd_v = [&](T lo, T hi) { return T(lo + (long long)(rng() % (unsigned long long)(hi - lo + 1))); };
    for (int it = 0; it < 300; it++) {
        int l = rng() % (n + 1), r = rng() % (n + 1);
        if (l > r) std::swap(l, r);
        T x = rnd_v(lo_v - 2, hi_v + 2), y = rnd_v(lo_v - 2, hi_v + 2);
        // count / sum
        int c1 = 0, c2 = 0;
        S s1 = S(), s2 = S();
        for (int i = l; i < r; i++) {
            if (a[i] < x) c1++, s1 += w[i];
            if (x <= a[i] && a[i] < y) c2++, s2 += w[i];
        }
        assert(wm.count(l, r, x) == c1 && wc.count(l, r, x) == c1);
        assert(wm.count(l, r, x, y) == c2 && wc.count(l, r, x, y) == c2);
        assert(wm.sum(l, r, x) == s1);
        assert(wm.sum(l, r, x, y) == s2);
        // prev / next
        std::optional<T> pv, nv;
        for (int i = l; i < r; i++) {
            if (a[i] < x && (!pv || *pv < a[i])) pv = a[i];
            if (x <= a[i] && (!nv || a[i] < *nv)) nv = a[i];
        }
        assert(wm.prev_value(l, r, x) == pv && wc.prev_value(l, r, x) == pv);
        assert(wm.next_value(l, r, x) == nv && wc.next_value(l, r, x) == nv);
        // kth / sum_smallest
        std::vector<int> idx(r - l);
        for (int i = l; i < r; i++) idx[i - l] = i;
        std::sort(idx.begin(), idx.end(), [&](int i, int j) { return a[i] != a[j] ? a[i] < a[j] : i < j; });
        S acc = S();
        assert(wm.sum_smallest(l, r, 0) == acc);
        for (int k = 0; k < r - l; k++) {
            assert(wm.kth_smallest(l, r, k) == a[idx[k]]);
            assert(wc.kth_smallest(l, r, k) == a[idx[k]]);
            assert(wm.kth_largest(l, r, k) == a[idx[r - l - 1 - k]]);
            acc += w[idx[k]];
            assert(wm.sum_smallest(l, r, k + 1) == acc);
        }
    }
}

int main() {
    std::mt19937 rng(12345);
    for (int it = 0; it < 200; it++) {
        int n = rng() % 30;
        int range = 1 + rng() % 20;
        std::vector<int> a(n);
        std::vector<long long> w(n);
        for (int i = 0; i < n; i++) {
            a[i] = int(rng() % range) - range / 2;
            w[i] = (long long)(rng() % 1000) - 500;
        }
        check<int, long long>(a, w, rng, -range / 2, range - range / 2);
    }
    // long long keys with large magnitude, few distinct values
    for (int it = 0; it < 50; it++) {
        int n = 1 + rng() % 40;
        std::vector<long long> a(n), w(n);
        std::vector<long long> pool = {-(1LL << 60), -7, 0, 3, 1LL << 62};
        int k = 1 + rng() % 5;
        for (int i = 0; i < n; i++) {
            a[i] = pool[rng() % k];
            w[i] = (long long)(rng() % 100);
        }
        wavelet_matrix<long long, long long> wm(a, w);
        for (int q = 0; q < 100; q++) {
            int l = rng() % (n + 1), r = rng() % (n + 1);
            if (l > r) std::swap(l, r);
            long long x = pool[rng() % 5] + (long long)(rng() % 3) - 1;
            int c = 0;
            long long s = 0;
            for (int i = l; i < r; i++) if (a[i] < x) c++, s += w[i];
            assert(wm.count(l, r, x) == c && wm.sum(l, r, x) == s);
            if (r > l) {
                std::vector<long long> b(a.begin() + l, a.begin() + r);
                std::sort(b.begin(), b.end());
                int kk = rng() % (r - l);
                assert(wm.kth_smallest(l, r, kk) == b[kk]);
            }
        }
    }
    // weights equal to the values
    {
        std::vector<int> a = {5, 1, 4, 1, 3};
        wavelet_matrix<int, long long> wm(a, std::vector<long long>(a.begin(), a.end()));
        assert(wm.sum_smallest(0, 5, 3) == 5);
        assert(wm.sum(0, 5, 1, 4) == 5);
        assert(wm.kth_smallest(0, 5, 2) == 3);
        assert(wm.kth_largest(0, 5, 0) == 5);
        assert(!wm.prev_value(0, 5, 1));
        assert(*wm.next_value(0, 5, 4) == 4);
        assert(!wm.next_value(0, 5, 6));
    }
    // empty
    {
        wavelet_matrix<int, long long> wm({}, {});
        assert(wm.size() == 0);
        assert(wm.count(0, 0, 10) == 0);
        assert(wm.sum(0, 0, 10) == 0);
        assert(wm.sum_smallest(0, 0, 0) == 0);
        assert(!wm.prev_value(0, 0, 10));
        wavelet_matrix<int> e;
        assert(e.size() == 0);
        wavelet_matrix<int, long long> ew;
        assert(ew.size() == 0 && ew.sum(0, 0, 0) == 0);
    }
    std::cout << "Hello World" << std::endl;
    return 0;
}
#line 1 "test/unit/wavelet_matrix.test.cpp"
#define PROBLEM "http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ITP1_1_A"
#include <algorithm>
#include <cassert>
#include <iostream>
#include <random>
#include <vector>

#line 3 "datastructure/wavelet_matrix.hpp"
#include <bit>
#line 5 "datastructure/wavelet_matrix.hpp"
#include <cstdint>
#include <optional>
#include <type_traits>
#line 9 "datastructure/wavelet_matrix.hpp"

template <class T, class S = void> struct wavelet_matrix {
private:
    using W = std::conditional_t<std::is_void_v<S>, int, S>;
    static constexpr bool weighted = !std::is_void_v<S>;

public:
    wavelet_matrix() {
        build({}, {});
    }
    explicit wavelet_matrix(const std::vector<T>& a)
        requires (!weighted)
    {
        build(a, {});
    }
    wavelet_matrix(const std::vector<T>& a, const std::vector<S>& w)
        requires weighted
    {
        assert(a.size() == w.size());
        build(a, w);
    }

    int size() const {
        return n;
    }
    // #{ i in [l, r) : a[i] < upper }
    int count(int l, int r, T upper) const {
        assert(0 <= l && l <= r && r <= n);
        return count_rank(l, r, rank_of(upper));
    }
    // #{ i in [l, r) : lower <= a[i] < upper }
    int count(int l, int r, T lower, T upper) const {
        assert(0 <= l && l <= r && r <= n);
        int lo = rank_of(lower), hi = rank_of(upper);
        if (lo >= hi) return 0;
        return count_rank(l, r, hi) - count_rank(l, r, lo);
    }
    // k-th smallest (0-indexed) of a[l..r)
    T kth_smallest(int l, int r, int k) const {
        assert(0 <= l && l <= r && r <= n);
        assert(0 <= k && k < r - l);
        int res = 0;
        for (int d = log - 1; d >= 0; d--) {
            int l0 = bv[d].rank0(l), r0 = bv[d].rank0(r);
            if (k < r0 - l0) {
                l = l0, r = r0;
            } else {
                k -= r0 - l0;
                res |= 1 << d;
                l += bv[d].zeros - l0, r += bv[d].zeros - r0;
            }
        }
        return key[res];
    }
    // k-th largest (0-indexed) of a[l..r)
    T kth_largest(int l, int r, int k) const {
        assert(0 <= l && l <= r && r <= n);
        assert(0 <= k && k < r - l);
        return kth_smallest(l, r, r - l - 1 - k);
    }
    // max { a[i] : i in [l, r), a[i] < upper }
    std::optional<T> prev_value(int l, int r, T upper) const {
        assert(0 <= l && l <= r && r <= n);
        int c = count_rank(l, r, rank_of(upper));
        if (c == 0) return std::nullopt;
        return kth_smallest(l, r, c - 1);
    }
    // min { a[i] : i in [l, r), lower <= a[i] }
    std::optional<T> next_value(int l, int r, T lower) const {
        assert(0 <= l && l <= r && r <= n);
        int c = count_rank(l, r, rank_of(lower));
        if (c == r - l) return std::nullopt;
        return kth_smallest(l, r, c);
    }
    // sum of w[i] over i in [l, r) with a[i] < upper
    S sum(int l, int r, T upper) const
        requires weighted
    {
        assert(0 <= l && l <= r && r <= n);
        return sum_rank(l, r, rank_of(upper));
    }
    // sum of w[i] over i in [l, r) with lower <= a[i] < upper
    S sum(int l, int r, T lower, T upper) const
        requires weighted
    {
        assert(0 <= l && l <= r && r <= n);
        int lo = rank_of(lower), hi = rank_of(upper);
        if (lo >= hi) return S();
        return sum_rank(l, r, hi) - sum_rank(l, r, lo);
    }
    // sum of w[i] over the k elements of a[l..r) with the smallest (a[i], i)
    S sum_smallest(int l, int r, int k) const
        requires weighted
    {
        assert(0 <= l && l <= r && r <= n);
        assert(0 <= k && k <= r - l);
        if (k == r - l) return cs[log][r] - cs[log][l];
        S res = S();
        for (int d = log - 1; d >= 0; d--) {
            int l0 = bv[d].rank0(l), r0 = bv[d].rank0(r);
            if (k < r0 - l0) {
                l = l0, r = r0;
            } else {
                k -= r0 - l0;
                res += cs[d][r0] - cs[d][l0];
                l += bv[d].zeros - l0, r += bv[d].zeros - r0;
            }
        }
        return res + cs[0][l + k] - cs[0][l];
    }

private:
    struct bit_vector {
        std::vector<uint64_t> bits;
        std::vector<int> cnt;
        int zeros;
        explicit bit_vector(int n) : bits(n / 64 + 1), cnt(n / 64 + 2), zeros(n) {
        }
        void set(int i) {
            bits[i >> 6] |= uint64_t(1) << (i & 63);
        }
        void build() {
            for (int i = 0; i < int(bits.size()); i++) {
                cnt[i + 1] = cnt[i] + std::popcount(bits[i]);
            }
            zeros -= cnt.back();
        }
        int rank1(int i) const {
            return cnt[i >> 6] +
                   std::popcount(bits[i >> 6] & ((uint64_t(1) << (i & 63)) - 1));
        }
        int rank0(int i) const {
            return i - rank1(i);
        }
    };

    int n, log;
    std::vector<T> key;
    std::vector<bit_vector> bv;
    // cs[d]: prefix sums of weights in the order at level d (cs[log]: original order)
    std::vector<std::vector<W>> cs;

    void build(const std::vector<T>& a, const std::vector<W>& w) {
        n = int(a.size());
        key = a;
        std::sort(key.begin(), key.end());
        key.erase(std::unique(key.begin(), key.end()), key.end());
        log = key.empty() ? 0 : std::bit_width(unsigned(key.size() - 1));
        std::vector<int> b(n), order(n);
        for (int i = 0; i < n; i++) {
            b[i] = int(std::lower_bound(key.begin(), key.end(), a[i]) - key.begin());
            order[i] = i;
        }
        bv.assign(log, bit_vector(n));
        if constexpr (weighted) {
            cs.assign(log + 1, std::vector<W>(n + 1));
            for (int i = 0; i < n; i++) cs[log][i + 1] = cs[log][i] + w[i];
        }
        std::vector<int> next(n);
        for (int d = log - 1; d >= 0; d--) {
            int z = 0;
            for (int i = 0; i < n; i++) {
                if (b[order[i]] >> d & 1) {
                    bv[d].set(i);
                } else {
                    next[z++] = order[i];
                }
            }
            for (int i = 0; i < n; i++) {
                if (b[order[i]] >> d & 1) next[z++] = order[i];
            }
            bv[d].build();
            std::swap(order, next);
            if constexpr (weighted) {
                for (int i = 0; i < n; i++) cs[d][i + 1] = cs[d][i] + w[order[i]];
            }
        }
    }
    // #{ v in key : v < x }
    int rank_of(const T& x) const {
        return int(std::lower_bound(key.begin(), key.end(), x) - key.begin());
    }
    // #{ i in [l, r) : b[i] < u }, 0 <= u <= |key|
    int count_rank(int l, int r, int u) const {
        if (u == 0) return 0;
        if (u >= int(key.size())) return r - l;
        int res = 0;
        for (int d = log - 1; d >= 0; d--) {
            int l0 = bv[d].rank0(l), r0 = bv[d].rank0(r);
            if (u >> d & 1) {
                res += r0 - l0;
                l += bv[d].zeros - l0, r += bv[d].zeros - r0;
            } else {
                l = l0, r = r0;
            }
        }
        return res;
    }
    S sum_rank(int l, int r, int u) const
        requires weighted
    {
        if (u == 0) return S();
        if (u >= int(key.size())) return cs[log][r] - cs[log][l];
        S res = S();
        for (int d = log - 1; d >= 0; d--) {
            int l0 = bv[d].rank0(l), r0 = bv[d].rank0(r);
            if (u >> d & 1) {
                res += cs[d][r0] - cs[d][l0];
                l += bv[d].zeros - l0, r += bv[d].zeros - r0;
            } else {
                l = l0, r = r0;
            }
        }
        return res;
    }
};
#line 9 "test/unit/wavelet_matrix.test.cpp"

template <class T, class S>
void check(const std::vector<T>& a, const std::vector<S>& w, std::mt19937& rng, T lo_v, T hi_v) {
    int n = int(a.size());
    wavelet_matrix<T, S> wm(a, w);
    wavelet_matrix<T> wc(a);
    assert(wm.size() == n && wc.size() == n);
    auto rnd_v = [&](T lo, T hi) { return T(lo + (long long)(rng() % (unsigned long long)(hi - lo + 1))); };
    for (int it = 0; it < 300; it++) {
        int l = rng() % (n + 1), r = rng() % (n + 1);
        if (l > r) std::swap(l, r);
        T x = rnd_v(lo_v - 2, hi_v + 2), y = rnd_v(lo_v - 2, hi_v + 2);
        // count / sum
        int c1 = 0, c2 = 0;
        S s1 = S(), s2 = S();
        for (int i = l; i < r; i++) {
            if (a[i] < x) c1++, s1 += w[i];
            if (x <= a[i] && a[i] < y) c2++, s2 += w[i];
        }
        assert(wm.count(l, r, x) == c1 && wc.count(l, r, x) == c1);
        assert(wm.count(l, r, x, y) == c2 && wc.count(l, r, x, y) == c2);
        assert(wm.sum(l, r, x) == s1);
        assert(wm.sum(l, r, x, y) == s2);
        // prev / next
        std::optional<T> pv, nv;
        for (int i = l; i < r; i++) {
            if (a[i] < x && (!pv || *pv < a[i])) pv = a[i];
            if (x <= a[i] && (!nv || a[i] < *nv)) nv = a[i];
        }
        assert(wm.prev_value(l, r, x) == pv && wc.prev_value(l, r, x) == pv);
        assert(wm.next_value(l, r, x) == nv && wc.next_value(l, r, x) == nv);
        // kth / sum_smallest
        std::vector<int> idx(r - l);
        for (int i = l; i < r; i++) idx[i - l] = i;
        std::sort(idx.begin(), idx.end(), [&](int i, int j) { return a[i] != a[j] ? a[i] < a[j] : i < j; });
        S acc = S();
        assert(wm.sum_smallest(l, r, 0) == acc);
        for (int k = 0; k < r - l; k++) {
            assert(wm.kth_smallest(l, r, k) == a[idx[k]]);
            assert(wc.kth_smallest(l, r, k) == a[idx[k]]);
            assert(wm.kth_largest(l, r, k) == a[idx[r - l - 1 - k]]);
            acc += w[idx[k]];
            assert(wm.sum_smallest(l, r, k + 1) == acc);
        }
    }
}

int main() {
    std::mt19937 rng(12345);
    for (int it = 0; it < 200; it++) {
        int n = rng() % 30;
        int range = 1 + rng() % 20;
        std::vector<int> a(n);
        std::vector<long long> w(n);
        for (int i = 0; i < n; i++) {
            a[i] = int(rng() % range) - range / 2;
            w[i] = (long long)(rng() % 1000) - 500;
        }
        check<int, long long>(a, w, rng, -range / 2, range - range / 2);
    }
    // long long keys with large magnitude, few distinct values
    for (int it = 0; it < 50; it++) {
        int n = 1 + rng() % 40;
        std::vector<long long> a(n), w(n);
        std::vector<long long> pool = {-(1LL << 60), -7, 0, 3, 1LL << 62};
        int k = 1 + rng() % 5;
        for (int i = 0; i < n; i++) {
            a[i] = pool[rng() % k];
            w[i] = (long long)(rng() % 100);
        }
        wavelet_matrix<long long, long long> wm(a, w);
        for (int q = 0; q < 100; q++) {
            int l = rng() % (n + 1), r = rng() % (n + 1);
            if (l > r) std::swap(l, r);
            long long x = pool[rng() % 5] + (long long)(rng() % 3) - 1;
            int c = 0;
            long long s = 0;
            for (int i = l; i < r; i++) if (a[i] < x) c++, s += w[i];
            assert(wm.count(l, r, x) == c && wm.sum(l, r, x) == s);
            if (r > l) {
                std::vector<long long> b(a.begin() + l, a.begin() + r);
                std::sort(b.begin(), b.end());
                int kk = rng() % (r - l);
                assert(wm.kth_smallest(l, r, kk) == b[kk]);
            }
        }
    }
    // weights equal to the values
    {
        std::vector<int> a = {5, 1, 4, 1, 3};
        wavelet_matrix<int, long long> wm(a, std::vector<long long>(a.begin(), a.end()));
        assert(wm.sum_smallest(0, 5, 3) == 5);
        assert(wm.sum(0, 5, 1, 4) == 5);
        assert(wm.kth_smallest(0, 5, 2) == 3);
        assert(wm.kth_largest(0, 5, 0) == 5);
        assert(!wm.prev_value(0, 5, 1));
        assert(*wm.next_value(0, 5, 4) == 4);
        assert(!wm.next_value(0, 5, 6));
    }
    // empty
    {
        wavelet_matrix<int, long long> wm({}, {});
        assert(wm.size() == 0);
        assert(wm.count(0, 0, 10) == 0);
        assert(wm.sum(0, 0, 10) == 0);
        assert(wm.sum_smallest(0, 0, 0) == 0);
        assert(!wm.prev_value(0, 0, 10));
        wavelet_matrix<int> e;
        assert(e.size() == 0);
        wavelet_matrix<int, long long> ew;
        assert(ew.size() == 0 && ew.sum(0, 0, 0) == 0);
    }
    std::cout << "Hello World" << std::endl;
    return 0;
}
Back to top page