跳转至

树状数组

树状数组上倍增

例题题意:给定一个长度为 $n$ 的序列 $a$,有 $q$ 次操作,每次操作:

  • 1 x y 将 $a_x$ 修改为 $y$。
  • 2 k 查询全局第 $k$ 小。

思路

使用权值树状数组维护每个数字的出现次数,即树状数组上 $t[i]$ 维护值域在 $[i-lowbit(i)+1,i]$ 的数字出现次数之和。

值域过大可以离散化。

对于全局第 $k$ 小,可以二分第 $k$ 小的值 $x$,然后通过

\[ \sum\limits_{i=1}^x cnt[i]<k \]

找到这个最大的 $x$,则 $x+1$ 就是全局第 $k$ 小。

这样做的时间复杂度为 $O(\log n^2)$。


考虑用倍增替代二分。

设 $x = 0$,$\mathrm{sum} = 0$,枚举 $i$ 从 $\log_2n$ 降为 $0$:

  • 查询权值数组中 $[x + 1 \ldots x + 2^i]$ 的区间和 $t$。
  • 如果 $\mathrm{sum} + t < k$,扩展成功,$x \gets x + 2^i$,$\mathrm{sum} \gets \mathrm{sum} + t$;否则扩展失败,不操作。

这样得到的 $x$ 是满足 $[1 \ldots x]$ 前缀和 $< k$ 的最大值,所以最终 $x + 1$ 就是答案。

看起来这种方法时间效率没有任何改善,但事实上,查询 $[x + 1 \ldots x + 2^i]$ 的区间和只需访问 $c[x + 2^i]$ 的值即可。

原因很简单,考虑 $\operatorname{lowbit}(x + 2^i)$,它一定是 $2^i$,因为 $x$ 之前只累加过 $2^j$ 满足 $j > i$。因此 $c[x + 2^i]$ 表示的区间就是 $[x + 1 \ldots x + 2^i]$。

如此一来,时间复杂度降低为 $\Theta(\log n)$。

参考代码
int kth(int k)
{
    int x = 0, sum = 0;
    for (int i = log2(n); i >= 0; i--)
    {
        int nxt = x + (1 << i);
        if (nxt <= n && sum + t[nxt] < k)
        {
            x = nxt;
            sum += t[nxt];
        }
    }
    return x + 1;
}

注意该方法只支持全局第 $k$ 小,因为只有查询是一段前缀才能利用树状数组的特性从前往后跳。否则不行


经典例题:P6619 [省选联考 2020 A/B 卷] 冰火战士

题意:有两类战士。

  • 冰系战士:场地温度必须 $\geq$ 其自身温度才能参赛;
  • 火系战士:场地温度必须 $\leq$ 其自身温度才能参赛。

确定场地温度后:

  • 冰系按温度升序(同温度能量大的优先)排队;
  • 火系按温度降序(同温度能量大的优先)排队。

比赛流程:双方队首战斗,能量小的耗尽并退出,能量大的剩余继续对抗对方下一个。如此往复,直到一方队列为空。 总消耗能量 = 双方所有战士在这场比赛中消耗的能量。

要求:在动态加入/撤销战士后,实时求当前情况下的最佳场地温度,使得总消耗能量最大,并输出该温度与对应的能量和。若无对战可能,则输出 Peace


解题思路:

关键点:由于战斗时,能量值小的一方会被消耗完毕。因此一旦确定了双方作战人员名单,则能量值等于冰火战斗人员能量之和的最小值。

且容易证明环境的最佳温度必然是某一个战斗人员的温度。(可以分类讨论得出)

假设环境温度为 $k$,则所求答案为以下和式

\[ \min(\sum\limits_{i\in \text{ice},x_i\leq k} y_i,\sum\limits_{i\in \text{fire},x_i\geq k} y_i) \]

其中 $\text{ice}$ 和 $\text{fire}$ 分别为冰系与火系。

因此不难看出冰系的参赛人员之和是一段前缀和,而火系的参赛人员之和是一段后缀和(也可以用参赛人员的能量值之和减去小于 $k$ 的能量值之和)。

假设温度 $k$ 确定后,就可以使用树状数组求解出冰系与火系的能量值之和了。


因此可以有如下暴力做法:

  • 由于温度值域过大,因此首先离散化所有温度存入 lsh 数组。
  • 维护一个 sum1 初始化为 $0$ 记录目前在场的火系人员的能量值之和,用于计算小于 $k$ 的能量值之和。
  • 枚举每一个温度,利用树状数组查询冰和火的能量值之和计算出最佳温度和最佳能量值即可。
参考代码
#include <bits/stdc++.h>
#define int long long
using namespace std;
constexpr int N = 2e6 + 5;

// 树状数组略

signed main()
{
    ios::sync_with_stdio(false), cin.tie(nullptr);
    int q;
    cin >> q;
    vector<array<int, 4>> ops(q);
    vector<int> lsh;
    for (int i = 0; i < q; i++)
    {
        int op;
        cin >> op;
        if (op == 1)
        {
            int t, x, y;
            cin >> t >> x >> y;
            ops[i] = {op, t, x, y};
            lsh.push_back(x);
        }
        else
        {
            int k;
            cin >> k;
            ops[i] = {op, k, 0, 0};
        }
    }
    sort(lsh.begin(), lsh.end());
    lsh.erase(unique(lsh.begin(), lsh.end()), lsh.end());
    auto get = [&](int x) -> int
    {
        return lower_bound(lsh.begin(), lsh.end(), x) - lsh.begin() + 1;
    };
    int m = lsh.size();
    tree t1(m), t0(m);
    int sum1 = 0;
    auto solve = [&]() -> pair<int, int>
    {
        int tem = -1, res = 0;
        for (int i = 0; i < m; i++) // 枚举所有温度 i + 1 求最佳答案
        {
            int ice = t0.query(i + 1); // 冰为小于等于温度 i + 1 的能量值之和
            int fire = sum1 - t1.query(i); // 火用在场人的和减去小于 $i+1$ 的
            if (ice == 0 || fire == 0) // 其中一方为 0 不合法
                continue;
            int tmp = min(ice, fire);
            if (tmp > res || (tmp == res && lsh[i] > tem))
            {
                res = tmp;
                tem = lsh[i];
            }
        }
        if (res == 0)
            return {-1, 0};
        return {tem, res};
    };

    int now = -1;
    for (int i = 0; i < q; i++)
    {
        auto [op, t, x, y] = ops[i];
        if (op == 1)
        {
            int id = get(x);
            if (t == 1)
            {
                t1.add(id, y);
                sum1 += y; // 火系在场人员登场
            }
            else
            {
                t0.add(id, y);
            }
        }
        else
        {
            auto it = ops[t - 1];
            int id = get(it[2]);
            if (it[1] == 1)
            {
                t1.add(id, -it[3]);
                sum1 -= it[3]; // 火系在场人员退场
            }
            else
            {
                t0.add(id, -it[3]);
            }
        }
        auto [tem, res] = solve();
        if (res == 0)
            cout << "Peace\n";
        else
            cout << tem << " " << res * 2 << "\n";
    }
    return 0;
}

二分优化

设 $ice(k)=\sum\limits_{i\in \text{ice},x_i\leq k} y_i$,$fire(k)=\sum\limits_{i\in \text{fire},x_i\geq k} y_i=sum1-\sum\limits_{i\in \text{fire},x_i< k} y_i$

不难发现 $ice(k)$ 随着 $k$ 的递增不下降,$fire(k)$ 随着 $k$ 的递增不上升。

因此两个函数可能存在一个交点,满足交点前 $ice(k) \leq fire(k)$,交点后 $ice(k) \geq fire(k)$。

由于函数存在连续相等的位置,因此无法三分找极值点,考虑使用二分。

  • 二分寻找最大的 $k_1$ 使得 $ice(k) < fire(k)$ 成立。
    • 则此时 $ice(k_1+1)\geq fire(k_1+1)$。
  • 但由于函数可能存在连续相等的段,因此在 $k_1+1$ 之后继续二分一个 $k_2$ 使得 $fire(k_2)=fire(k_1+1)$。
  • 那么答案要么是 $k_1$ 要么是 $k_2$。

那么就有一个 $O(q\log{n}^2)$ 的做法了。

关键代码
auto solve = [&]() -> pair<int, int>
{
    int tem = -1, res = 0;
    int l = 0, r = m - 1, k1 = -1;
    // 二分最大的 k1 使得 ice(k) < fire(k)
    while (l <= r)
    {
        int mid = l + r >> 1;
        int ice = t0.query(mid + 1), fire = sum1 - t1.query(mid);
        if (ice < fire) k1 = mid, l = mid + 1;
        else r = mid - 1;
    }
    // k1 存在先求一次答案
    if (k1 != -1)
    {
        int ice = t0.query(k1 + 1), fire = sum1 - t1.query(k1);
        if (ice > 0 && fire > 0)
        {
            int tmp = min(ice, fire);
            if (tmp > res || (tmp == res && lsh[k1] > tem))
            {
                res = tmp;
                tem = lsh[k1];
            }
        }
    }
    // k1 + 1 < m 才存在 k2
    if (k1 != -1 && k1 + 1 < m)
    {
        // 在 k1 + 1 到 m - 1 中二分找 k2 使得 fire(k2) = fire(k1 + 1)
        int l = k1 + 1, r = m - 1;
        int k2 = 0, tag = sum1 - t1.query(k1 + 1);
        while (l <= r)
        {
            int mid = l + r >> 1;
            int fire = sum1 - t1.query(mid);
            if (fire == tag) k2 = mid, l = mid + 1;
            else if (fire < tag) r = mid - 1;
            else l = mid + 1;
        }
        if (k2 != -1)
        {
            int ice = t0.query(k2 + 1), fire = sum1 - t1.query(k2);
            int tmp = min(ice, fire);
            if (tmp > res || (tmp == res && lsh[k2] > tem))
            {
                res = tmp;
                tem = lsh[k2];
            }
        }
    }
    return {tem, res};
};

树状数组上倍增

找 $k_1$ 的过程实际是找到一个最大的前缀,使得 $ice(k)<fire(k)$。那么这个过程可以使用树状数组上倍增(或二分)来实现。

具体来说,贪心的从大步往小步去跳,设当前位置为 $k_1$ 初始化为 $0$。

  • 记下一个位置 $pos=k_1+(2^i)$。
    • 易知:$\text{lowbit(pos)}=2^i$,且树状数组上 $t[pos]$ 维护的区间恰好就是 $[k_1+1,k_1+2^i]$。
  • 这样就利用起了树状数组的特性,就可以倍增找到最大的 $k_1$ 使得 $ice(k)<fire(k)$。

接下来找火系最远位置 $k_2$,注意到火系能量的计算方式为 $sum1-\sum\limits_{x_i<k} y_i$。这里是小于 $k$ 取不到相等,这不利于我们树状数组上倍增,因为 $\text{lowbit(pos)}=2^i$,而 $\text{lowbit(pos - 1)}$ 显然不一定是 $2^i-1$。

可以维护火系树状数组时,让每个战斗人员的温度值偏移一位(都增加 $1$)这样查询火系的战斗值之和也变为了 $sum1-\sum\limits_{x_i\leq k} y_i$

然后同样的树状数组上倍增维护找到 $k_2$ 即可。

整体时间复杂度:$O(q\log{n})$。

关键代码
auto solve = [&]() -> pair<int, int>
{
    int tem = -1, res = 0;
    int k1 = 0, ice = 0, fire = sum1;
    for (int i = 20; i >= 0; i--)
    {
        int nxt = k1 + (1 << i);
        if (nxt <= m)
        {
            int nice = ice + t0.t[nxt];
            int nfire = fire - t1.t[nxt];
            if (nice < nfire)
            {
                k1 = nxt;
                ice = nice, fire = nfire;
            }
        }
    }
    if (ice > 0 && fire > 0)
    {
        tem = lsh[k1], res = min(ice, fire);
    }
    if (k1 < m)
    {
        int k2 = 0, tag = sum1 - t1.query(k1 + 1);
        int sum = sum1;
        for (int i = 20; i >= 0; i--)
        {
            int nxt = k2 + (1 << i);
            if (nxt <= m)
            {
                int nsum = sum - t1.t[nxt];
                if (nsum >= tag)
                {
                    k2 = nxt;
                    sum = nsum;
                }
            }
        }
        ice = t0.query(k2), fire = sum1 - t1.query(k2);
        int tmp = min(ice, fire);
        if (tmp > res || (tmp == res && lsh[k2] > tem))
        {
            tem = lsh[k2];
            res = tmp;
        }
    }
    return {tem, res};
};