DataStructure/OrderedSet.h

• View this file on GitHub
• Last update: 2022-01-17 03:27:26+08:00
• Link: https://cses.fi/problemset/task/1076
• Link: https://open.kattis.com/problems/lazylearner

Code

// Copied from chemthan
// 
// Usage:
// 
// OrderedTree<int> set;
// set.init(MAX_VALUE);  // Init set to store 0..n-1
// set.upd(x, x, +1);    // normally can just upd x = val (val is just for mapping of values)
// set.find_by_order(k)  // k from 1
//
// Tested:
// - https://open.kattis.com/problems/lazylearner
// - https://cses.fi/problemset/task/1076

template<class num_t>
struct OrderedTree {
    int n;
    vector<num_t> a;
    num_t tot;
    vector<num_t> x;

    void init(int _n) {
        n = _n;
        a.resize(n);
        for (int i = 0; i < n; i++) a[i] = 0;
        tot = 0;
        x.resize(n + 1);
        for (int i = 0; i <= n; i++) x[i] = 0;
    }

    void upd(int p, num_t val, num_t cnt) {
        assert(0 <= p && p < n);
        a[p] = val, tot += cnt;
        for (p++; p <= n; p += p & -p) {
            x[p] += cnt;
        }
    }

    // 1 <= k <= tot  (tot = number of elements)
    num_t find_by_order(num_t k) {
        assert(0 <= k && k <= tot);
        if (!k) return -1;
        int p = 0;
        for (int i = __lg(n); i >= 0; i--) {
            if (p + (1 << i) <= n && x[p + (1 << i)] < k) {
                k -= x[p + (1 << i)];
                p += 1 << i;
            }
        }
        return a[p];
    }
};

#line 1 "DataStructure/OrderedSet.h"
// Copied from chemthan
// 
// Usage:
// 
// OrderedTree<int> set;
// set.init(MAX_VALUE);  // Init set to store 0..n-1
// set.upd(x, x, +1);    // normally can just upd x = val (val is just for mapping of values)
// set.find_by_order(k)  // k from 1
//
// Tested:
// - https://open.kattis.com/problems/lazylearner
// - https://cses.fi/problemset/task/1076

template<class num_t>
struct OrderedTree {
    int n;
    vector<num_t> a;
    num_t tot;
    vector<num_t> x;

    void init(int _n) {
        n = _n;
        a.resize(n);
        for (int i = 0; i < n; i++) a[i] = 0;
        tot = 0;
        x.resize(n + 1);
        for (int i = 0; i <= n; i++) x[i] = 0;
    }

    void upd(int p, num_t val, num_t cnt) {
        assert(0 <= p && p < n);
        a[p] = val, tot += cnt;
        for (p++; p <= n; p += p & -p) {
            x[p] += cnt;
        }
    }

    // 1 <= k <= tot  (tot = number of elements)
    num_t find_by_order(num_t k) {
        assert(0 <= k && k <= tot);
        if (!k) return -1;
        int p = 0;
        for (int i = __lg(n); i >= 0; i--) {
            if (p + (1 << i) <= n && x[p + (1 << i)] < k) {
                k -= x[p + (1 << i)];
                p += 1 << i;
            }
        }
        return a[p];
    }
};

Back to top page