用 SBT 解决第 k 小的数

 一 问题描述

有 N 个数字 1, 2, 3, …, N ，在第 i 轮游戏中，伊萨找出第 ki 小的数并把它拿走，求解 M 轮游戏后伊萨拿走的数字总和。

二 输入和输出

1 输入

第 1 行输入一个整数 T ，表示测试用例的数量。每个测试用例的第 1 行都包含两个整数 N 和 M ，N 表示数字的个数，M 表示游戏的轮数。接下来的一行包含 M 个数，第 i 个数 ki 表示伊萨在第 i  轮取出第 ki 小的数。其中：1≤T≤128；1≤M≤N≤262144；1≤ki≤N-i+1。

2 输出

对每个测试用例，先输出用例号，再输出总和。

三 输入和输出样例

1 输入样例

3

3 2

1 1

10 3

3 9 1

2 输出样例

Case 1:3

Case 2:14

四 分析和设计

本问题求多轮第 k 小的元素之和，可以采用 SBT 解决。

五 代码

package com.platform.modules.alg.alglib.hdu4217;

public class Hdu4217 {
    public String output = "";
    private int maxn = 100005;
    int n, cnt; // 结点数，结点存储下标累计
    node tr[] = new node[maxn];

    public String cal(String input) {
        int T, n, m, k, tmp, cas = 0;
        long ans;

        String[] line = input.split("\n");
        T = Integer.parseInt(line[0]);

        int count = 1;
        while (T-- > 0) {
            ans = cnt = 0;
            root = new nodePtr();
            String[] num = line[count++].split(" ");
            n = Integer.parseInt(num[0]);
            for (int i = 1; i <= n; i++)
                insert(root, i);
            String[] ki = line[count++].split(" ");
            for (int i = 0; i < ki.length; i++) {
                k = Integer.parseInt(ki[i]);
                ans += tmp = get_kth(root, k);
                remove(root, tmp);
            }
            output += "Case " + (++cas) + ":" + ans + "\n";
        }

        return output;
    }

    public Hdu4217() {
        for (int i = 0; i < tr.length; i++) {
            tr[i] = new node();
        }
    }

    void R_rotate(nodePtr x) {
        int y = tr[x.ptr].lc.ptr;
        tr[x.ptr].lc.ptr = tr[y].rc.ptr;
        tr[y].rc.ptr = x.ptr;
        tr[y].size = tr[x.ptr].size;
        tr[x.ptr].size = tr[tr[x.ptr].lc.ptr].size + tr[tr[x.ptr].rc.ptr].size + 1;
        x.ptr = y;
    }

    void L_rotate(nodePtr x) {
        int y = tr[x.ptr].rc.ptr;
        tr[x.ptr].rc.ptr = tr[y].lc.ptr;
        tr[y].lc.ptr = x.ptr;
        tr[y].size = tr[x.ptr].size;
        tr[x.ptr].size = tr[tr[x.ptr].lc.ptr].size + tr[tr[x.ptr].rc.ptr].size + 1;
        x.ptr = y;
    }

    void maintain(nodePtr p, boolean flag) {
        if (p.ptr == 0) return;
        if (!flag) {
            if (tr[tr[tr[p.ptr].lc.ptr].lc.ptr].size > tr[tr[p.ptr].rc.ptr].size) // LL
                R_rotate(p);
            else if (tr[tr[tr[p.ptr].lc.ptr].rc.ptr].size > tr[tr[p.ptr].rc.ptr].size) { // LR
                L_rotate(tr[p.ptr].lc);
                R_rotate(p);
            } else return;
        } else {
            if (tr[tr[tr[p.ptr].rc.ptr].rc.ptr].size > tr[tr[p.ptr].lc.ptr].size)//RR
                L_rotate(p);
            else if (tr[tr[tr[p.ptr].rc.ptr].lc.ptr].size > tr[tr[p.ptr].lc.ptr].size) {//RL
                R_rotate(tr[p.ptr].rc);
                L_rotate(p);
            } else return;
        }
        maintain(tr[p.ptr].lc, false);
        maintain(tr[p.ptr].rc, true);
        maintain(p, false);
        maintain(p, true);
    }

    void insert(nodePtr p, int val) {
        if (p.ptr == 0) {
            p.ptr = ++cnt;
            tr[p.ptr].lc.ptr = tr[p.ptr].rc.ptr = 0;
            tr[p.ptr].size = 1;
            tr[p.ptr].val = val;
        } else {
            tr[p.ptr].size++;
            if (val < tr[p.ptr].val) insert(tr[p.ptr].lc, val);
            else insert(tr[p.ptr].rc, val);
            maintain(p, val >= tr[p.ptr].val);
        }
    }

    int remove(nodePtr p, int val) {
        tr[p.ptr].size--;
        if ((tr[p.ptr].val == val) || (tr[p.ptr].val > val && 0 == tr[p.ptr].lc.ptr) || (tr[p.ptr].val < val && 0 == tr[p.ptr].rc.ptr)) {
            int tmp = tr[p.ptr].val;
            if (0 == tr[p.ptr].lc.ptr || 0 == tr[p.ptr].rc.ptr)//有一个儿子为空，直接用儿子代替
                p.ptr = tr[p.ptr].lc.ptr + tr[p.ptr].rc.ptr;
            else // 找前驱，左子树最右节点
                tr[p.ptr].val = remove(tr[p.ptr].lc, tr[p.ptr].val + 1);
            return tmp;
        } else if (val < tr[p.ptr].val) return remove(tr[p.ptr].lc, val);
        else return remove(tr[p.ptr].rc, val);
    }

    // 求第 k 小数
    int get_kth(nodePtr p, int k) { // 求第 k 小数
        int s = tr[tr[p.ptr].lc.ptr].size + 1;
        if (s == k) return tr[p.ptr].val;
        else if (s < k) return get_kth(tr[p.ptr].rc, k - s);
        else return get_kth(tr[p.ptr].lc, k);
    }
    
    // 树根
    nodePtr root = new nodePtr();
}

class node {
    nodePtr lc = new nodePtr(); // 左孩子
    nodePtr rc = new nodePtr(); // 右孩子
    int val; // 值
    int size; // 子树大小
}

class nodePtr {
    int ptr;
}

六 测试