[CF1399E2] Weights Division (hard version)

Description

给定一棵有根树，每条边有一个 w 和一个 c，设这棵树的权值为所有叶子结点到根结点的路径上的 w 的和的总和。每次操作可以选择一条边，将其 w 变为 w/2（向下取整），花费 c 的代价。

Solution

先预处理出每条边的次数 (d)

如果只有一种代价，我们只需要将所有边作为 ((w,d)) 扔进大根堆中，以 ((w-[frac w 2])d) 作为关键字排序。

有两种代价时，考虑维护两个优先队列，一个存 (c=1) 的，另一个存 (c=2) 的。比较行动哪个队列更优。

#include <bits/stdc++.h>
using namespace std;

#define int long long
const int N = 1000005;

struct edge
{
    int v, w, c;
};

vector<edge> g[N];
int vis[N], d[N], c[N], w[N], n, m, s, sum, ans;

struct Item
{
    int w, d;
    int op1() const
    {
        return (w - w / 2) * d;
    }
    int op2() const
    {
        return (w - w / 2 / 2) * d;
    }
    bool operator<(const Item &b) const
    {
        return op1() < b.op1();
    }
};

priority_queue<Item> que[3];

void clear()
{
    for (int i = 1; i <= n; i++)
    {
        g[i].clear();
        vis[i] = d[i] = c[i] = w[i] = 0;
    }
    while (que[1].size())
        que[1].pop();
    while (que[2].size())
        que[2].pop();
    sum = ans = 0;
}

void dfs(int p)
{
    vis[p] = 1;
    int flag = 0;
    for (edge e : g[p])
    {
        int q = e.v;
        if (!vis[q])
        {
            flag = 1;
            c[q] = e.c;
            w[q] = e.w;
            dfs(q);
            d[p] += d[q];
        }
    }
    if (flag == 0)
        d[p] = 1;
}

void solve()
{
    cin >> n >> s;
    for (int i = 1; i < n; i++)
    {
        int t1, t2, t3, t4;
        cin >> t1 >> t2 >> t3 >> t4;
        g[t1].push_back({t2, t3, t4});
        g[t2].push_back({t1, t3, t4});
    }
    dfs(1);

    /* for (int i = 1; i <= n; i++)
    {
        cout << "i=" << i << " c=" << c[i] << " w=" << w[i] << " d=" << d[i] << endl;
    } */

    for (int i = 2; i <= n; i++)
    {
        que[c[i]].push({w[i], d[i]});
        sum += w[i] * d[i];
    }

    while (sum > s)
    {
        if (que[1].size())
        {
            int tt = que[1].top().op1();
            // cout << "tt " << tt << "  " << endl;
            if (sum - tt <= s)
            {
                ans++;
                break;
            }
        }

        int t1 = 0, t2 = 0, t3 = 0;
        if (que[1].size())
            t1 = que[1].top().op2();
        if (que[1].size() > 1)
        {
            auto tmp = que[1].top();
            t2 = que[1].top().op1();
            que[1].pop();
            t2 += que[1].top().op1();
            que[1].push(tmp);
        }
        if (que[2].size()) t3 = que[2].top().op1();
        // cout << "::: " << t1 << " " << t2 << " " << t3 << endl;
        if (t1 >= t2 && t1 >= t3)
        {
            auto tmp1 = que[1].top();
            sum -= tmp1.op1();
            que[1].pop();
            tmp1.w /= 2;
            que[1].push(tmp1);
            ans += 1;
        }
        else if (t2 >= t1 && t2 >= t3)
        {
            auto tmp1 = que[1].top();
            sum -= tmp1.op1();
            que[1].pop();
            tmp1.w /= 2;
            que[1].push(tmp1);
            ans += 1;
        }
        else
        {
            auto tmp = que[2].top();
            sum -= tmp.op1();
            que[2].pop();
            tmp.w /= 2;
            que[2].push(tmp);
            ans += 2;
        }
    }
    cout << ans << endl;
    clear();
}

signed main()
{
    ios::sync_with_stdio(false);

    int t;
    cin >> t;
    while (t--)
    {
        solve();
    }
}