题目出处HDU 4217 Data Structure?
题意大致如下(由deepseek总结)
- 初始有一个包含 1 到 N 的整数集合（所有数互不相同，按自然顺序）。
- 进行 K 轮操作，第 i 轮给出一个正整数 Ki，要求在当前剩余的数字中，找到第 Ki 小的数，将它从集合中移除，并累加到答案总和中。
- 输出所有被移除的数的总和。
首先看一下为什么这是树状数组的题目
这道题目里我们需要支持两个操作：
- 查询第$k_i $小的数
- 删除这个数
先考虑暴力求解
- 用数组维护当前剩余的数，每次查询\(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;
}