よくあるデータ構造問題！！　めっちゃ色んな解法がある！

問題へのリンク

問題概要

長さ $N$ の整数列 $A_{0}, A_{1}, \dots, A_{N-1}$ と整数 $M, K$ が与えられる (0-indexed で表している)。

各 $i = 0, 1, \dots, N-M$ に対して、次の問題に答えてください。

$M$ 個の整数 $A_{i}, A_{i+1}, \dots, A_{i+M-1}$ を小さい順に並び替えたときの先頭 $K$ 個の総和を求めよ。

制約

$1 \le K \le M \le N \le 2 \times 10^{5}$
$1 \le A_{i} \le 10^{9}$

考えたこと

要は次のクエリを高速に処理できるデータ構造があればよい。

値 $x$ を挿入する
値 $x$ を削除する (複数個ある場合は 1 個だけ削除する)
小さい方から数えて $k$ 番目の値を求める

このデータ構造を用いて、「値 $A_{i}$ を push する」「値 $A_{i}$ を pop する」という処理 (後述) をそれぞれ実装しておくことで、今回の問題は次のように解ける。

最初に $A_{0}, A_{1}, \dots, A_{M-1}$ を push する
各に対して
- データ構造中の小さい順に $K$ 個の総和を出力する
- 値 $A_{i+M}$ を push する
- 値 $A_{i}$ を pop する

具体的には、以下のように実現できる。なお、変数 sum を、データ構造中の小さい順に $K$ 個の総和 ( $K$ 個未満の場合はすべての総和) とする。

値 $A_{i}$ を push

データ構造に値 $A_{i}$ を挿入して、sum += A[i] とする。

そしてこの操作によって「小さい順に $K$ 個以内」から追い出される数 v があれば、sum -= v とする。

値 $A_{i}$ を pop

データ構造から値 $A_{i}$ を削除する。この値が「小さい順に $K$ 個以内」に含まれるならば、sum -= A[i] とする。

さらにこの操作によって「小さい順に $K$ 個以内」に新たに加わる数 v があれば、sum += v とする。

データ構造をどうするか

以上の push / pop においては、冒頭で述べた通り、次の 3 種類のクエリを高速に実行できる必要がある。

値 $x$ を挿入する
値 $x$ を削除する (複数個ある場合は 1 個だけ削除する)
小さい方から数えて $k$ 番目の値を求める

そのようなデータ構造は無数に考えられる。

multiset
削除機能を備えた priority_queue
BIT 上二分探索の機能を備えた BIT
セグメント木
Binary Trie
Wavelet Matrix

解法 (1)：multiset (C++)

最も高等知識を必要としない解法。

$K$ 番目の値を取得するためには、有名な方法がある。それは

小さい順に $K$ 個を管理するための multiset：left
それ以上の値を管理するための multiset：right

とを 2 つ持つ戦法だ。値を挿入・削除するときには、うまく left と right とのバランスを取るように実装すればよい。なお、multiset $S$ 中の最小の要素は *S.begin()、最大の要素は *S.rbegin() で取得できる。

注意点として、C++ の multiset には有名な罠がある。multiset $S$ から値 $x$ を削除したいときに、

S.erase(x);

と書いてしまうと、 $S$ に含まれるすべての $x$ が削除されてしまうのだ。今回は $x$ を 1 個だけ削除したいので注意が必要だ。

正しくは、削除したいイテレータ (ここでは it とする) を求めて S.erase(it) のようにする必要がある。具体的には、関数 find() を用いて、

S.erase(S.find(x));

のように書けばよい！

全体として計算量は $O(N \log N)$ となる。

コード (解法 (1))

AC コードを開く

#include <bits/stdc++.h>
using namespace std;

int main() {
    // 入力
    long long N, M, K;
    cin >> N >> M >> K;
    vector<long long> A(N);
    for (int i = 0; i < N; ++i) cin >> A[i];

    // 小さい順に K 個の総和
    long long sum = 0;

    // multiset
    multiset<long long> left, right;

    auto push = [&](long long x) -> void {
        // とりあえず x を left に挿入する
        left.insert(x);
        sum += x;

        // left のサイズが K を超えるなら left の最大値を right に移す
        if (left.size() > K) {
            long long y = *left.rbegin();
            right.insert(y);
            sum -= y;
            left.erase(left.find(y));
        }
    };
    auto pop = [&](long long x) -> void {
        // とりあえず x を削除する
        if (left.count(x)) {
            left.erase(left.find(x));
            sum -= x;
        } else {
            right.erase(right.find(x));
        }

        // left のサイズが K 未満になるなら right の最小値を left に移す
        if (left.size() < K) {
            long long y = *right.begin();
            left.insert(y);
            sum += y;
            right.erase(right.find(y));
        }
    };

    for (int i = 0; i < M; ++i) push(A[i]);
    for (int i = 0; i < N - M + 1; ++i) {
        cout << sum << " ";
        if (i+M < N) push(A[i+M]), pop(A[i]);
    }
    cout << endl;
}

解法 (2)：削除機能を備えた priority_queue

multiset と同様に、priority_queue を用いて $K$ 番目の値を管理するためには、

小さい順に $K$ 個を管理するための priority_queue (最大の値を取得できるようにする)：left
それ以上の値を管理するための priority_queue (最小の値を取得できるようにする)：right

を用意する方法が有名だ。

しかし、priority_queue で実現するにあたって一つ障壁がある。それは「priority_queue から値 $x$ を削除したいときどうするのか」という部分だ。

通常の priority_queue que では「最大の値」「最小の値」しか削除できない。この問題を乗り越えるために、削除したい値を予約しておく priority_queue del を別に用意しておく。そして、que から「最大の値」を取得するときに、次の疑似コードのように「取得した値が del で予約された値と一致するかどうかを確認し、一致するならば棄却する」ようにすればよい。

// 削除機能を持たせた priority_queue
priority_queue<int> que, del;

// 挿入
void push(int x) { que.push(x); }

// 削除
void pop(int x) { del.push(x); }

// 最大の値を取得 (ここでは最大の値の削除はしない)
int get() {
    while (!del.empty() && que.top() == del.top()) {
        que.pop();
        del.pop();
    }
    return que.top();
}

この方法を活用して、priority_queue を合計 4 個用意することで、この問題は $O(N \log N)$ で解ける。

コード (解法 (2))

AC コードを開く

#include <bits/stdc++.h>
using namespace std;

int main() {
    // 入力
    long long N, M, K;
    cin >> N >> M >> K;
    vector<long long> A(N);
    for (int i = 0; i < N; ++i) cin >> A[i];

    // 小さい順に K 個の総和
    long long sum = 0;

    // proiroty_queue
    priority_queue<long long> left, dleft;
    priority_queue<long long, vector<long long>, greater<long long>> right, dright;

    // 境界値を取得する
    auto get_left = [&]() -> long long {
        while (!dleft.empty() && left.top() == dleft.top()) {
            left.pop(); dleft.pop();
        }
        return left.top();
    };
    auto get_right = [&]() -> long long {
        while (!dright.empty() && right.top() == dright.top()) {
            right.pop(); dright.pop();
        }
        return right.top();
    };

    // push と pop
    auto push = [&](long long x) -> void {
        // とりあえず x を left に挿入する
        left.push(x);
        sum += x;

        // left のサイズが K を超えるなら left の最大値を right に移す
        if (left.size() - dleft.size() > K) {
            long long y = get_left();
            sum -= y;
            left.pop();
            right.push(y);
        }
    };
    auto pop = [&](long long x) -> void {
        // とりあえず x を削除する
        if (x <= get_left()) {
            dleft.push(x);
            sum -= x;
        } else {
            dright.push(x);
        }

        // left のサイズが K 未満になるなら right の最小値を left に移す
        if (left.size() - dleft.size() < K) {
            long long y = get_right();
            sum += y;
            right.pop();
            left.push(y);
        }
    };

    for (int i = 0; i < M; ++i) push(A[i]);
    for (int i = 0; i < N - M + 1; ++i) {
        cout << sum << " ";
        if (i+M < N) push(A[i+M]), pop(A[i]);
    }
    cout << endl;
}

解法 (3)：BIT 上二分探索の機能を備えた BIT

ここから先は高級なデータ構造を使う！

「挿入」「削除」「 $K$ 番目を取得」クエリを処理するのに BIT (Binary Indexed Tree) を使うとよいケースも多い。ただし、挿入する値は $0$ 以上 $10^{6}$ (程度) 以下の整数でなければならない。BIT 内部で用意する配列 dat に対して、値 x を挿入するときは dat[x] にアクセスするためだ。

そこで、今回の問題では座標圧縮をする。配列 $A$ の各値を小さい順に番号付ける。それにより、 $A$ の各値は $0, 1, \dots, N-1$ を並び替えたものになる (今回はタイは潰さない)。これで BIT が使えるようになる。

なお、BIT を用いて具体的に「挿入」「削除」「 $K$ 番目を取得」を実行する方法については、たとえば次の資料などを参考に。 $K$ 番目の値を取得する部分については、「BIT 上の二分探索」がキーワード。

algo-logic.info

コード (解法 (3))

小さい順に $K$ 個の総和を求める部分では、これまでのように変数 sum を用いてもよいが、下コードのように、総和を求めるための専用の BIT 変数をもう 1 つ用意する方が簡便だと思う！

AC コードを開く

#include <bits/stdc++.h>
using namespace std;

// get(k): binary search on BIT
template <class Abel> struct BIT {
    Abel UNITY_SUM = 0;
    vector<Abel> dat;
    
    // [0, n)
    BIT(int n, Abel unity = 0) : UNITY_SUM(unity), dat(n, unity) { }
    void init(int n) {
        dat.assign(n, UNITY_SUM);
    }
    
    // a is 0-indexed
    inline void add(int a, Abel x) {
        for (int i = a; i < (int)dat.size(); i |= i + 1)
            dat[i] = dat[i] + x;
    }
    
    // [0, a), a is 0-indexed
    inline Abel sum(int a) {
        Abel res = UNITY_SUM;
        for (int i = a - 1; i >= 0; i = (i & (i + 1)) - 1)
            res = res + dat[i];
        return res;
    }
    
    // [a, b), a and b are 0-indexed
    inline Abel sum(int a, int b) {
        return sum(b) - sum(a);
    }

    // k-th number (k is 0-indexed)
    int get(long long k) {
        ++k;
        int res = 0;
        int N = 1;
        while (N < (int)dat.size()) N *= 2;
        for (int i = N / 2; i > 0; i /= 2) {
            if (res + i - 1 < (int)dat.size() && dat[res + i - 1] < k) {
                k = k - dat[res + i - 1];
                res = res + i;
            }
        }
        return res;
    }
    
    // debug
    void print() {
        for (int i = 0; i < (int)dat.size(); ++i)
            cout << sum(i, i + 1) << ",";
        cout << endl;
    }
};

int main() {
    // 入力
    long long N, M, K;
    cin >> N >> M >> K;
    vector<long long> A(N);
    for (int i = 0; i < N; ++i) cin >> A[i];

    // 座標圧縮を準備
    vector<pair<long long,int>> comp(N);
    for (int i = 0; i < N; ++i) comp[i] = {A[i], i};
    sort(comp.begin(), comp.end());
    vector<int> order(N);
    for (int i = 0; i < N; ++i) order[comp[i].second] = i;

    // BIT
    BIT<long long> bit(N + 1), sum(N + 1);

    // push と pop
    auto push = [&](long long x, int id) -> void {
        bit.add(id, 1);
        sum.add(id, x);
    };
    auto pop = [&](long long x, int id) -> void {
        bit.add(id, -1);
        sum.add(id, -x);
    };
    auto get = [&]() -> long long {
        int kthid = bit.get(K-1);
        return sum.sum(kthid+1);
    };

    for (int i = 0; i < M; ++i) push(A[i], order[i]);
    for (int i = 0; i < N - M + 1; ++i) {
        cout << get() << " ";
        if (i+M < N) {
            push(A[i+M], order[i+M]); 
            pop(A[i], order[i]);
        }
    }
    cout << endl;
}

解法 (4)：セグメント木

BIT で解けるならセグメント木でも解ける。解法 (3) で準備した 2 つの BIT を相当して、

各ノードの値：(挿入された値の個数, 挿入された値の総和)
ノード間の演算：(挿入された値の個数の和, 挿入された値の総和の和)

によって定義されたセグメント木を用意する。

「小さい順に $K$ 個の総和」を求めるクエリを処理するためには次のようにする。

区間 $\lbrack 0, r)$ に対するセグメント木の取得値の第一要素が $K$ 以下になるような最大の $r$ を求める (セグメント木上の二分探索で求められる)
1 で求めた $r$ に対して、区間 $\lbrack 0, r)$ に対するセグメントの取得値の第二要素を答える

コード (解法 (4))

ACL を用いた。ACL には、セグメント木上の二分探索を実行する関数 max_right() も用意されている。

AC コードを開く

#include <bits/stdc++.h>
#include <atcoder/segtree>
using namespace std;
using namespace atcoder;

// セグメント木の設定
using Monoid = pair<long long, long long>;
Monoid op(Monoid a, Monoid b) {
    return {a.first + b.first, a.second + b.second};
};
Monoid e(){return {0, 0};}

int main() {
    // 入力
    long long N, M, K;
    cin >> N >> M >> K;
    vector<long long> A(N);
    for (int i = 0; i < N; ++i) cin >> A[i];

    // 座標圧縮を準備
    vector<pair<long long,int>> comp(N);
    for (int i = 0; i < N; ++i) comp[i] = {A[i], i};
    sort(comp.begin(), comp.end());
    vector<int> order(N);
    for (int i = 0; i < N; ++i) order[comp[i].second] = i;

    // セグメント木
    segtree<Monoid, op, e> seg(N);

    // セグメント木上の二分探索を実行するための関数
    auto f = [&](Monoid x) -> bool { return x.first <= K; };

    // push と pop
    auto push = [&](long long x, int id) -> void {
        seg.set(id, Monoid(1, x));
    };
    auto pop = [&](long long x, int id) -> void {
        seg.set(id, e());
    };
    auto get = [&]() -> long long {
        int r = seg.max_right(0, f);
        return seg.prod(0, r).second;
    };

    for (int i = 0; i < M; ++i) push(A[i], order[i]);
    for (int i = 0; i < N - M + 1; ++i) {
        cout << get() << " ";
        if (i+M < N) {
            push(A[i+M], order[i+M]); 
            pop(A[i], order[i]);
        }
    }
    cout << endl;
}

解法 (5)：Binary Trie

「挿入」「削除」「 $K$ 番目の値を取得」ができるデータ構造はまだまだある。

挿入する値が非負整数値ならば、Binary Trie も有力。詳細はここでは省略。

コード (解法 (5))

AC コードを開く

#include <bits/stdc++.h>
using namespace std;

// Binary Trie
template<typename INT, size_t MAX_DIGIT> struct BinaryTrie {
    struct Node {
        size_t count;
        Node *prev, *left, *right;
        Node(Node *prev) : count(0), prev(prev), left(nullptr), right(nullptr) {}
    };
    INT lazy;
    Node *root;

    // constructor
    BinaryTrie() : lazy(0), root(emplace(nullptr)) {}
    inline size_t get_count(Node *v) const { return v ? v->count : 0; }
    inline size_t size() const { return get_count(root); }

    // add and get value of Node
    inline void add(INT val) {
        lazy ^= val;
    }
    inline INT get(Node *v) {
        if (!v) return -1;
        INT res = 0;
        for (int i = 0; i < MAX_DIGIT; ++i) {
            if (v == v->prev->right)
                res |= INT(1)<<i;
            v = v->prev;
        }
        return res ^ lazy;
    }

    // find Node* whose value is val
    Node* find(INT val) {
        INT nval = val ^ lazy;
        Node *v = root;
        for (int i = MAX_DIGIT-1; i >= 0; --i) {
            bool flag = (nval >> i) & 1;
            if (flag) v = v->right;
            else v = v->left;
            if (!v) return v;
        }
        return v;
    }

    // insert
    inline Node* emplace(Node *prev) {
        return new Node(prev);
    }
    void insert(INT val, size_t k = 1) {
        INT nval = val ^ lazy;
        Node *v = root;
        for (int i = MAX_DIGIT-1; i >= 0; --i) {
            bool flag = (nval >> i) & 1;
            if (flag && !v->right) v->right = emplace(v);
            if (!flag && !v->left) v->left = emplace(v);
            if (flag) v = v->right;
            else v = v->left;
        }
        v->count += k;
        while ((v = v->prev)) v->count = get_count(v->left) + get_count(v->right);
    }
    
    // erase
    Node* clear(Node *v) {
        if (!v || get_count(v)) return v;
        delete(v);
        return nullptr;
    }
    bool erase(Node *v, size_t k = 1) {
        if (!v) return false;
        v->count -= k;
        while ((v = v->prev)) {
            v->left = clear(v->left);
            v->right = clear(v->right);
            v->count = get_count(v->left) + get_count(v->right);
        }
        return true;
    }
    bool erase(INT val) {
        auto v = find(val);
        return erase(v);
    }

    // max (with xor-addition of val) and min (with xor-addition of val)
    Node* get_max(INT val = 0) {
        INT nval = val ^ lazy;
        Node* v = root;
        for (int i = MAX_DIGIT-1; i >= 0; --i) {
            bool flag = (nval >> i) & 1;
            if (!v->right) v = v->left;
            else if (!v->left) v = v->right;
            else if (flag) v = v->left;
            else v = v->right;
        }
        return v;
    }
    Node* get_min(INT val = 0) {
        return get_max(~val & ((INT(1)<<MAX_DIGIT)-1));
    }
   
    // lower_bound, upper_bound
    Node* get_cur_node(Node *v, int i) {
        if (!v) return v;
        Node *left = v->left, *right = v->right;
        if ((lazy >> i) & 1) swap(left, right);
        if (left) return get_cur_node(left, i+1);
        else if (right) return get_cur_node(right, i+1);
        return v;
    }
    Node* get_next_node(Node *v, int i) {
        if (!v->prev) return nullptr;
        Node *left = v->prev->left, *right = v->prev->right;
        if ((lazy >> (i+1)) & 1) swap(left, right);
        if (v == left && right) return get_cur_node(right, i);
        else return get_next_node(v->prev, i+1);
    }
    Node* lower_bound(INT val) {
        INT nval = val ^ lazy;
        Node *v = root;
        for (int i = MAX_DIGIT-1; i >= 0; --i) {
            bool flag = (nval >> i) & 1;
            if (flag && v->right) v = v->right;
            else if (!flag && v->left) v = v->left;
            else if ((val >> i) & 1) return get_next_node(v, i);
            else return get_cur_node(v, i);
        }
        return v;
    }
    Node* upper_bound(INT val) {
        return lower_bound(val + 1);
    }
    size_t order_of_val(INT val) {
        Node *v = root;
        size_t res = 0;
        for (int i = MAX_DIGIT-1; i >= 0; --i) {
            Node *left = v->left, *right = v->right;
            if ((lazy >> i) & 1) swap(left, right);
            bool flag = (val >> i) & 1;
            if (flag) {
                res += get_count(left);
                v = right;
            }
            else v = left;
        }
        return res;
    }

    // k-th, k is 0-indexed
    Node* get_kth(size_t k, INT val = 0) {
        Node *v = root;
        if (get_count(v) <= k) return nullptr;
        for (int i = MAX_DIGIT-1; i >= 0; --i) {
            bool flag = (lazy >> i) & 1;
            Node *left = (flag ? v->right : v->left);
            Node *right = (flag ? v->left : v->right);
            if (get_count(left) <= k) k -= get_count(left), v = right;
            else v = left;
        }
        return v;
    }

    // debug
    void print(Node *v, string prefix = "") {
        if (!v) return;
        cout << prefix << ": " << v->count << endl;
        print(v->left, prefix + "0");
        print(v->right, prefix + "1");
    }
    void print() {
        print(root);
    }
    vector<INT> eval(Node *v, int digit) const {
        vector<INT> res;
        if (!v) return res;
        if (!v->left && !v->right) {
            for (int i = 0; i < get_count(v); ++i) res.push_back(0);
            return res;
        }
        const auto& left = eval(v->left, digit-1);
        const auto& right = eval(v->right, digit-1);
        for (auto val : left) res.push_back(val);
        for (auto val : right) res.push_back(val + (INT(1)<<digit));
        return res;
    }
    vector<INT> eval() const {
        auto res = eval(root, MAX_DIGIT-1);
        for (auto &val : res) val ^= lazy;
        return res;
    }
    friend ostream& operator << (ostream &os,
                                 const BinaryTrie<INT, MAX_DIGIT> &bt) {
        auto res = bt.eval();
        for (auto val : res) os << val << " ";
        return os;
    }
};

int main() {
    // 入力
    long long N, M, K;
    cin >> N >> M >> K;
    vector<long long> A(N);
    for (int i = 0; i < N; ++i) cin >> A[i];

    // 小さい順に K 個の総和
    long long sum = 0;

    // Binary Trie
    BinaryTrie<int, 30> left, right;

    // push と pop
    auto push = [&](long long x) -> void {
        // とりあえず x を left に挿入する
        left.insert(x);
        sum += x;

        // left のサイズが K を超えるなら left の最大値を right に移す
        if (left.size() > K) {
            long long y = left.get(left.get_max());
            sum -= y;
            left.erase(y);
            right.insert(y);
        }
    };
    auto pop = [&](long long x) -> void {
        // とりあえず x を削除する
        if (x <= left.get(left.get_max())) {
            left.erase(x);
            sum -= x;
        } else {
            right.erase(x);
        }

        // left のサイズが K 未満になるなら right の最小値を left に移す
        if (left.size() < K) {
            long long y = right.get(right.get_min());
            sum += y;
            left.insert(y);
            right.erase(y);
        }
    };

    for (int i = 0; i < M; ++i) push(A[i]);
    for (int i = 0; i < N - M + 1; ++i) {
        cout << sum << " ";
        if (i+M < N) push(A[i+M]), pop(A[i]);
    }
    cout << endl;
}

解法 (6)：Wavelet Matrix

ユーザ解説 by rsk0315 にて、上位互換問題が Wavelet Matrix で解けることが指摘されている。

コード (解法 (6))

AC コードを開く

#include <bits/stdc++.h>
using namespace std;


// Bit Vector
struct BitRank {
    // block: bit vector
    // count: the number of 1 within each block
    vector<unsigned long long> block;
    vector<unsigned int> count;
    
    // constructor
    BitRank() {}
    void resize(const unsigned int num) {
        block.resize(((num + 1) >> 6) + 1, 0);
        count.resize(block.size(), 0);
    }
    
    // set val(0 or 1) onto i-th bit
    void set(const unsigned int i, const unsigned long long val) {
        block[i >> 6] |= (val << (i & 63));
    }
    void build() {
        for (unsigned int i = 1; i < block.size(); i++) {
            count[i] = count[i - 1] + __builtin_popcountll(block[i - 1]);
        }
    }
    
    // the number of 1 in [0, i)
    unsigned int rank1(const unsigned int i) const {
        return count[i >> 6] +
        __builtin_popcountll(block[i >> 6] & ((1ULL << (i & 63)) - 1ULL));
    }
    // the number of 1 in [i, j)
    unsigned int rank1(const unsigned int i, const unsigned int j) const {
        return rank1(j) - rank1(i);
    }
    // the number of 0 in [0, i)
    unsigned int rank0(const unsigned int i) const {
        return i - rank1(i);
    }
    // the number of 0 in [i, j)
    unsigned int rank0(const unsigned int i, const unsigned int j) const {
        return rank0(j) - rank0(i);
    }
};

// Wavelet Matrix
template<class T> struct WaveletMatrix {
    // inner data
    int height;
    vector<BitRank> bv;
    vector<int> pos;
    vector<vector<long long>> rui;

    // constructor (sigma: the number of characters)
    WaveletMatrix() {}
    WaveletMatrix(vector<T> vec) :
        WaveletMatrix(vec, *max_element(vec.begin(), vec.end()) + 1) {}
    WaveletMatrix(vector<T> vec, T sigma) {
        init(vec, sigma);
    }
    void init(vector<T> &vec, T sigma) {
        height = (sigma == 1) ? 1 : (64 - __builtin_clzll(sigma - 1));
        bv.resize(height), pos.resize(height);
        vector<T> A = vec;
        rui.resize(height + 1);
        for (int i = 0; i < height; ++i) {
            bv[i].resize(vec.size());
            for (int j = 0; j < vec.size(); ++j) {
                bv[i].set(j, get(vec[j], height - i - 1));
            }
            bv[i].build();
            auto it = stable_partition(vec.begin(), vec.end(), [&](int c) {
                return !get(c, height - i - 1);
            });
            pos[i] = it - vec.begin();
        }
        for (int i = 0; i <= height; ++i) {
            rui[i].resize((int)A.size() + 1);
            for (int j = 1; j <= (int)A.size(); ++j) {
                rui[i][j] = rui[i][j - 1] + A[j - 1];
            }
            if (i == height) break;
            stable_partition(A.begin(), A.end(), [&](int c) {
                return !get(c, height - i - 1);
            });
        }
    }
    
    // the i-th bit of "val" (0 or 1)
    int get(const T val, const int i) {
        return val >> i & 1;
    }
    
    // the number of "val" in [l, r)
    int rank(const T val, const int l, const int r) {
        return rank(val, r) - rank(val, l);
    }
    
    // the number of "val" in [0, i)
    int rank(T val, int i) {
        int p = 0;
        for (int j = 0; j < height; ++j) {
            if (get(val, height - j - 1)) {
                p = pos[j] + bv[j].rank1(p);
                i = pos[j] + bv[j].rank1(i);
            } else {
                p = bv[j].rank0(p);
                i = bv[j].rank0(i);
            }
        }
        return i - p;
    }
    
    // the k-th (0-indexed) smallest value in [l, r)
    T k_th_smallest(int k, int l, int r) {
        T res = 0;
        for (int i = 0; i < height; ++i) {
            const int j = bv[i].rank0(l, r);
            if (j > k){
                l = bv[i].rank0(l);
                r = bv[i].rank0(r);
            } else {
                l = pos[i] + bv[i].rank1(l);
                r = pos[i] + bv[i].rank1(r);
                k -= j;
                res |= (1LL << (height - i - 1));
            }
        }
        return res;
    }
    
    // the k-th (0-indexed) largest value in [l, r)
    T k_th_largest(int l, int r, int k) {
       return k_th_smallest(l, r, r - l - k - 1);
    }
    
    // the sum of the top-k sum in [l, r)
    T top_k_sum(int k, int l, int r) {
        if (l == r) return 0;
        T res = 0, val = 0;
        for (int i = 0; i < height; ++i) {
            const int j = bv[i].rank0(l, r);
            if (j > k) {
                l = bv[i].rank0(l);
                r = bv[i].rank0(r);
            } else {
                int l2 = bv[i].rank0(l);
                int r2 = bv[i].rank0(r);
                res += rui[i + 1][r2] - rui[i + 1][l2];
                l = pos[i] + bv[i].rank1(l);
                r = pos[i] + bv[i].rank1(r);
                k -= j;
                val |= (1LL << (height - i - 1));
            }
        }
        res += (long long)val * k;
        return res;
    }
    
    // the number of value between [loewr, upper) in [l, r)
    int range_freq(const int i, const int j, const T lower, const T upper,
                   const int l, const int r, const int x) {
        if (i == j || r <= lower || upper <= l) return 0;
        int mid = (l + r) >> 1;
        if (lower <= l && r <= upper) {
            return j - i;
        } else {
            T left = range_freq(bv[x].rank0(i), bv[x].rank0(j), lower, upper, l, mid, x + 1);
            T right = range_freq(pos[x] + bv[x].rank1(i), pos[x] + bv[x].rank1(j),
                                 lower, upper, mid, r, x + 1);
            return left + right;
        }
    }
    int range_freq(const int l, const int r, const T lower, const T upper){
        return range_freq(l, r, lower, upper, 0, 1LL << height, 0);
    }
    
    // the minmum value between [lower, upper) in [l, r)
    int range_min(const int i, const int j, const T lower, const T upper,
                  const int l, const int r, const int x, const int val){
        if (i == j || r <= lower || upper <= l) return -1;
        if (r - l == 1) return val;
        int mid = (l + r) >> 1;
        int res = range_min(bv[x].rank0(i), bv[x].rank0(j), lower, upper, l, mid, x + 1, val);
        if (res < 0) {
            return range_min(pos[x] + bv[x].rank1(i), pos[x] + bv[x].rank1(j),
                             lower, upper, mid, r, x + 1, val + (1LL << (height - x - 1)));
        } else return res;
    }
    int range_min(int l, int r, T lower, T upper) {
        return range_min(l, r, lower, upper, 0, 1LL << height, 0, 0);
    }
    
    // the max value (< val) in [l, r)
    T prev_value(int l, int r, T val) {
        int num = range_freq(l, r, 0, val);
        if (num == 0) return T(-1);
        else return k_th_smallest(l, r, num - 1);
    }
    
    // the min value (>= val) in [l, r)
    T next_value(int l, int r, T val) {
        int num = range_freq(l, r, 0, val);
        if (num == r - l) return T(-1);
        else return k_th_smallest(l, r, num);
    }
};

//
template<class T> struct OrthogonalRangeCount {
    // inner data
    using ptt = pair<T, T>;
    const int sz;
    vector<T> X, Y;
    WaveletMatrix<T> wm;
    
    // constructor
    OrthogonalRangeCount(vector<ptt> candidate) : sz((int)candidate.size()), X(sz), Y(sz) {
        sort(candidate.begin(), candidate.end());
        vector<int> vec(sz);
        for (int i = 0; i < sz; ++i) {
            X[i] = candidate[i].first, Y[i] = candidate[i].second;
        }
        sort(Y.begin(), Y.end());
        Y.erase(unique(Y.begin(), Y.end()), Y.end());
        for (int i = 0; i < sz; ++i) {
            vec[i] = lower_bound(Y.begin(), Y.end(), candidate[i].second) - Y.begin();
        }
        wm.init(vec, Y.size());
    }
    
    // the number of points in [lx, rx) × [ly, ry)
    int query(const T lx, const T ly, const T rx, const T ry){
        const int lxid = lower_bound(X.begin(), X.end(), lx) - X.begin();
        const int rxid = lower_bound(X.begin(), X.end(), rx) - X.begin();
        const int lyid = lower_bound(Y.begin(), Y.end(), ly) - Y.begin();
        const int ryid = lower_bound(Y.begin(), Y.end(), ry) - Y.begin();
        if (lxid >= rxid || lyid >= ryid) return 0;
        return wm.range_freq(lxid, rxid, lyid, ryid);
    }
};

void ABC_281_E() {
    int N, M, K;
    cin >> N >> M >> K;
    vector<long long> A(N);
    for (int i = 0; i < N; ++i) cin >> A[i];
    
    WaveletMatrix wm(A);
    for (int i = 0; i < N - M + 1; ++i) {
        cout << wm.top_k_sum(K, i, i + M) << " ";
    }
    cout << endl;
}

int main() {
    ABC_281_E();
}

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 281 E - Least Elements (水色, 500 点)

問題概要

制約

考えたこと

値 $A_{i}$ を push

値 $A_{i}$ を pop

データ構造をどうするか

解法 (1)：multiset (C++)

コード (解法 (1))

解法 (2)：削除機能を備えた priority_queue

コード (解法 (2))

解法 (3)：BIT 上二分探索の機能を備えた BIT

コード (解法 (3))

解法 (4)：セグメント木

コード (解法 (4))

解法 (5)：Binary Trie

コード (解法 (5))

解法 (6)：Wavelet Matrix

コード (解法 (6))

問題概要

制約

考えたこと

値 を push

値 を pop

データ構造をどうするか

解法 (1)：multiset (C++)

コード (解法 (1))

解法 (2)：削除機能を備えた priority_queue

コード (解法 (2))

解法 (3)：BIT 上二分探索の機能を備えた BIT

コード (解法 (3))

解法 (4)：セグメント木

コード (解法 (4))

解法 (5)：Binary Trie

コード (解法 (5))

解法 (6)：Wavelet Matrix

コード (解法 (6))

値 $A_{i}$ を push

値 $A_{i}$ を pop