HDUOJ 4217 Data Structure?

HDUOJ 4217 Data Structure?

题目链接

Problem Description

Data structure is one of the basic skills for Computer Science students, which is a particular way of storing and organizing data in a computer so that it can be used efficiently. Today let me introduce a data-structure-like problem for you.
Original, there are N numbers, namely 1, 2, 3…N. Each round, iSea find out the Ki-th smallest number and take it away, your task is reporting him the total sum of the numbers he has taken away.

Input

The first line contains a single integer T, indicating the number of test cases.
Each test case includes two integers N, K, K indicates the round numbers. Then a line with K numbers following, indicating in i (1-based) round, iSea take away the Ki-th smallest away.

Technical Specification

1. 1 <= T <= 128
2. 1 <= K <= N <= 262 144
3. 1 <= Ki <= N - i + 1

Output

For each test case, output the case number first, then the sum.

Sample Input

2
3 2
1 1
10 3
3 9 1

Sample Output

Case 1: 3
Case 2: 14

典型的权值线段树查询第 $k$ 小，每次查询一个后把那个位置的数置为 0 即可，注意答案会爆 $int$，AC代码如下：

#include "bits/stdc++.h"

using namespace std;
typedef long long ll;
const int N = 5e5 + 5;
int n, t, q, k;

class WeightedSegmentTree {
private:
    vector<ll> tree, add;
public:
    void init(int n) {
        tree.resize(n << 2);
        add.resize(n << 2);
    }

    void pushup(int i) {
        tree[i] = tree[i << 1] + tree[i << 1 | 1];
    }

    void pushdown(int i, int l, int r) {
        int mid = l + (r - l) / 2;
        tree[i << 1] += ll(mid - l + 1) * add[i];
        tree[i << 1 | 1] += ll(r - mid) * add[i];
        add[i << 1] = add[i << 1 | 1] = add[i];
        add[i] = 0;
    }

    void build(int i, int l, int r) {
        add[i] = 0;
        if (l == r) {
            tree[i] = 1;
            return;
        }
        int mid = l + (r - l) / 2;
        build(i << 1, l, mid);
        build(i << 1 | 1, mid + 1, r);
        pushup(i);
    }

    void update(int i, int l, int r, int m, int n, ll k) {//区间加k
        if (m <= l && r <= n) {
            add[i] = add[i] + k;
            tree[i] = tree[i] + (r - l + 1LL) * k;
            return;
        }
        pushdown(i, l, r);
        int mid = l + (r - l) / 2;
        if (m <= mid) update(i << 1, l, mid, m, n, k);
        if (n > mid) update(i << 1 | 1, mid + 1, r, m, n, k);
        pushup(i);
    }

    ll query1(int i, int l, int r, int m, int n) {//查询区间[l,r]中的数出现的次数和
        if (m <= l && r <= n) return tree[i];
        ll ans = 0;
        int mid = l + (r - l) / 2;
        pushdown(i, l, r);
        if (m <= mid) ans = ans + query1(i << 1, l, mid, m, n);
        if (n > mid) ans = ans + query1(i << 1 | 1, mid + 1, r, m, n);
        return ans;
    }

    ll query2(int i, int l, int r, int k) {//查询区间[l,r]内的第k小，第k小对应第n-k+1大
        if (l == r) return l;
        int mid = l + (r - l) / 2;
        if (tree[i << 1] >= k) return query2(i << 1, l, mid, k);
        else return query2(i << 1 | 1, mid + 1, r, k - tree[i << 1]);
    }
};


int main() {
    scanf("%d", &t);
    for (int i = 1; i <= t; i++) {
        scanf("%d%d", &n, &q);
        WeightedSegmentTree tree = WeightedSegmentTree();
        tree.init(n);
        tree.build(1, 1, n);
        ll ans = 0;
        while (q--) {
            scanf("%d", &k);
            ll num = tree.query2(1, 1, n, k);
            ans += num;
            tree.update(1, 1, n, num, num, -1);
        }
        printf("Case %d: %lld\n", i, ans);
    }
    return 0;
}