题目链接
EPIC Institute of Technology Round Summer 2024 (Div. 1 + Div. 2) E. Wonderful Tree!
思路
题目要求令所有的 a v ≤ ∑ u ∈ s o n ( v ) a u a_{v} \le \sum_{u \in son(v)}a_{u} av≤∑u∈son(v)au。我们令 b v = ∑ u ∈ s o n ( v ) a u − a v b_{v} = \sum_{u \in son(v)}a_{u} - a_{v} bv=∑u∈son(v)au−av，则问题转化为了使所有的 b v ≥ 0 b_{v} \ge 0 bv≥0的最少操作数。
对于任意节点 v v v，对其子树中除自己的任意一个节点 u u u，使 b u = b u − 1 b_{u} = b_{u} - 1 bu=bu−1， b v = b v + 1 b_{v} = b_{v} + 1 bv=bv+1，操作的代价是 u u u到 v v v的距离。
对于每一个 b < 0 b < 0 b<0的节点，我们只要贪心的先选深度小且 b b b为正值的节点，用它的 b b b值来填充自己的。叶子节点的 b b b值设置为正无穷即可。
代码
#include <bits/stdc++.h>
using namespace std;
#define int long long
typedef pair<int, int> pii;
const int N = 5e3 + 5, M = 1e4 + 5;
const int mod = 1e9 + 7;
const int inf = 0x3f3f3f3f3f3f3f3f;
int n, ans;
int a[N], b[N], depth[N];
vector<int>mp[N];
void dfs1(int u, int deep)
{
b[u] = -a[u];
depth[u] = deep;
for (int j : mp[u])
{
dfs1(j, deep + 1);
b[u] += a[j];
}
if (!mp[u].size()) b[u] = inf;
}
vector<int> dfs2(int u)
{
vector<int>res;
for (int j : mp[u])
{
vector<int>v = dfs2(j);
for (int val : v)
res.push_back(val);
}
sort(res.begin(), res.end(), [&](int x, int y) {return depth[x] > depth[y];});
if (b[u] > 0) res.push_back(u);
else if (b[u] < 0)
{
int idx = res.size() - 1;
while (b[u] < 0)
{
int v = res[idx];
int sum = min(-b[u], b[v]);
b[v] -= sum, b[u] += sum;
ans += sum * (depth[v] - depth[u]);
if (b[v] == 0) idx--, res.pop_back();
}
}
return res;
}
void solve()
{
cin >> n;
ans = 0;
for (int i = 1; i <= n; i++)
{
mp[i].clear();
}
for (int i = 1; i <= n; i++)
{
cin >> a[i];
}
for (int i = 2, p; i <= n; i++)
{
cin >> p;
mp[p].push_back(i);
}
dfs1(1, 1);
dfs2(1);
cout << ans << endl;
}
signed main()
{
ios::sync_with_stdio(false);
cin.tie(0), cout.tie(0);
int test = 1;
cin >> test;
for (int i = 1; i <= test; i++)
{
solve();
}
return 0;
}