树状数组训练题 - HDU 4217 Data Structure?

题目出处HDU 4217 Data Structure?

题意大致如下(由deepseek总结)

• 初始有一个包含 1 到 N 的整数集合（所有数互不相同，按自然顺序）。
• 进行 K 轮操作，第 i 轮给出一个正整数 Ki，要求在当前剩余的数字中，找到第 Ki 小的数，将它从集合中移除，并累加到答案总和中。
• 输出所有被移除的数的总和。

首先看一下为什么这是树状数组的题目

这道题目里我们需要支持两个操作：

1. 查询第$k_i $小的数
2. 删除这个数

先考虑暴力求解

• 用数组维护当前剩余的数，每次查询\(O(n)\)
• 如果单纯数组维护，删除\(O(n)\)，如果链表维护，删除\(O(1)\)
  总计时间复杂度\(O(n * k)\)，在题目2.6e6的数据规模下很容易TLE

换个思路

这个时候我们需要换个思路，不能单纯地去维护当前剩余的数了。
由于我们要求的是第\(k_i\)小的数，而且给定的初始数组是天然有序的（类似std::iota生成的数组），下标与值相等
所以，若记当前的下标为x，数组为tr我们可以考虑维护tr[x]的值为小于等于x的值的个数。

再考虑使用数据结构优化

这个时候我们看到了单点查询和单点修改，只有这两个需求，而且需要降低复杂度，很容易就能想到树状数组了。

现在来看实现：

把模板中原来的前缀和与查询改成以下代码，用于查找第k小的数

i32 kth(i32 k)
{
    i32 idx = 0;
    for (i32 bit = 1 << (31 - __builtin_clz(n)); bit; bit >>= 1)
    {
        i32 nxt = idx + bit;
        if (nxt <= n && tr[nxt] < k)
        {
            k -= tr[nxt];
            idx = nxt;
        }
    }
    return idx + 1;
}

同时，初始化时，依次把所有的值都增加1，可以实现把tr数组维护成tr[i] = i - 1

for (i32 i = 1; i <= n; ++ i)
    bit.add(i, 1);

最后贴完整的ac代码

#include <bits/stdc++.h>
using namespace std;
using i32 = int;
using i64 = long long;

struct BIT
{
    i32 n; vector<i64> tr;

    BIT(i32 n_)
    {
        n = n_;
        tr.assign(n + 1, 0);
    }

    i64 lowbit(i32 x)
    {
        return x & -x;
    }

    void add(i32 idx, i32 delta)
    {
        while (idx <= n)
        {
            tr[idx] += delta;
            idx += lowbit(idx);
        }
    }

    i32 kth(i32 k)
    {
        i32 idx = 0;
        for (i32 bit = 1 << (31 - __builtin_clz(n)); bit; bit >>= 1)
        {
            i32 nxt = idx + bit;
            if (nxt <= n && tr[nxt] < k)
            {
                k -= tr[nxt];
                idx = nxt;
            }
        }
        return idx + 1;
    }

};

void solve(i32 testcase)
{
    i32 n, k; cin >> n >> k;
    BIT bit = BIT(n);
    for (i32 i = 1; i <= n; ++ i)
        bit.add(i, 1);
    
    i64 ans = 0;
    while (k -- )
    {
        i32 x; cin >> x;
        i32 val = bit.kth(x);
        ans += val;
        bit.add(val, -1);
    }

    cout << "Case " << testcase << ": " << ans << '\n'; 
}

i32 main()
{
    ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
    i32 t; cin >> t;
    for (i32 i = 1; i <= t; i ++ ) solve(i);
    return 0;
}