Codechef Tree Balancing (TREEBAL)
题目链接
Codechef Tree Balancing (TREEBAL)
题目大意
给定一棵大小为 \(n\) 的树，根节点为节点 \(1\)， 第 \(i\) 条边的长度为 \(w_i\)。
你可以进行若干次修改，每次选择一条边 \(i\)，给 \(w_i\) 加上 \(1\) 或减去 \(1\)，代价为 \(c_i\)，注意 \(w_i\) 可以修改成负数。
你的目标是让所有叶子到根节点的距离相同，求达到目标的最小代价，以及任意一组最优方案。
\(1\leq n\leq 2\times 10^5,\;1\leq w_i,c_i\leq 10^6\)
思路
朴素的树 \(dp\) 想法显然是设 \(dp_{u,i}\) 表示子树 \(u\) 内的所有叶子深度（在算上整棵子树和 \(u\) 的父边时）为 \(i\) 时的最小代价，转移为 \(dp_{u,i}=\min_{l}\{|w_e - l|+\sum dp_{v,i-l}\}\)。
注意到绝对值函数是凸函数，而 \(dp\) 转移的过程是一系列凸函数的叠加，所以 \(dp_{u,i}\) 是一个关于 \(i\) 的分段线性凸函数，且每次只会增加 \(O(1)\) 个段，于是考虑转而维护凸包。
维护凸包上的所有拐点的横坐标，每个拐点对应两侧斜率变化了 \(1\)，从而一些坐标会出现很多次，于是用 set 维护二元组 \((x,num)\) 表示坐标和出现次数，另外需要知道凸包的斜率和具体数值，维护凸包最右侧的射线的表达式 \(y=ax+b\)，即记录 \((a,b)\) 。
叶子：\(dp_u:(w_e,2c_e)\)，\(y=c_ex-c_ew_e\)
合并儿子信息：启发式合并所有 set，将 \((a,b)\) 对应相加。
增加 \(u\) 的父边 \(e\)：\(e\) 的意义在于当凸包的斜率绝对值 \(>c_e\) 时，这些地方的深度变化用修改 \(e\) 的长度会更优，于是 删去凸包两侧斜率在 \([-c_e,c_e]\) 之外的拐点，在删除右侧拐点 \(x\) （一个）时，\((a,b)\rightarrow (a-1,b+x)\)。同时所 有拐点向右平移 \(w_e\)，此时 \((a,b)\rightarrow (a,b-aw_e)\) 。
在求答案时，一个个删去斜率 \(>0\) 的拐点，最终得到最底部斜率为 \(0\) 的线段端点，以及表达式 \(y=0x+b\)，于是答案为 \(b\)，具体深度任取线段上一值即可。
对于求方案，注意到一条边 \(e\) 只在原斜率绝对值 \(>c_e\) 时才发挥作用，所以我们维护好每个点的凸包的最左和最右两个拐点 \(l_i,r_i\)，以及两侧射线的斜率 \(mn_i,mx_i\)。设我们当前求的是子树 \(u\) 以及父边 \(e\) 内，叶子深度 \(=x\) 的方案，若 $ x>r_i$ 且 \(mx_i=c_e\)，则 \(w_e\) 要加上 \(x-r_i\)，若 \(x<l_i\) 且 \(mn_i=-c_e\)，则 \(w_e\) 要减去 \(l_i-x\)，否则 \(e\) 对最优方案没啥贡献，不变。然后子树内 \(x\rightarrow x-w_e\) 递归求解即可。
时间复杂度 \(O(n\log^2n)\) 。
Code
#include<iostream>
#include<cstring>
#include<set>
#define mem(a,b) memset(a, b, sizeof(a))
#define rep(i,a,b) for(int i = (a); i <= (b); i++)
#define per(i,b,a) for(int i = (b); i >= (a); i--)
#define N 200021
#define ll long long
#define PLL pair<ll, ll>
#define fr first
#define sc second
using namespace std;
int n;
int head[N], nxt[2*N], to[2*N];
ll w[2*N], c[2*N], len[2*N];
int cnt;
struct Hull{
set<PLL> s;
ll a, b, shift, num;
void adjust(ll c, bool id){
if(id) while(a-num < -c){
ll need = -c-a + num;
PLL p = *s.begin();
num -= min(p.sc, need), s.erase(s.begin());
if(p.sc > need) s.insert({p.fr, p.sc-need});
}
while(a > c){
PLL p = *(--s.end());
ll del = min(a-c, p.sc);
a -= del, b += del * (p.fr+shift), num -= del, s.erase(--s.end());
if(p.sc > del) s.insert({p.fr, p.sc-del});
}
}
} hull[N];
ll lb[N], rb[N], mn[N], mx[N];
void init(){
rep(i,1,n) head[i] = -1; cnt = -1;
rep(i,1,n) hull[i].s.clear(), hull[i].a = hull[i].b = hull[i].shift = hull[i].num = 0;
}
void add_e(int a, int b, int w, int c, bool id){
nxt[++cnt] = head[a], head[a] = cnt, to[cnt] = b;
::w[cnt] = w, ::c[cnt] = c;
if(id) add_e(b, a, w, c, 0);
}
void merge(Hull &a, Hull &b){
if(a.s.size() < b.s.size()) swap(a, b);
for(PLL p : b.s){
ll x = p.fr + b.shift - a.shift;
auto it = a.s.lower_bound({x, 0});
if(it != a.s.end() && it->fr == x) p.sc += it->sc, a.s.erase(it);
a.s.insert({x, p.sc});
}
a.num += b.num;
a.a += b.a, a.b += b.b;
}
void dfs(int x, int e){
bool leaf = true;
for(int i = head[x]; ~i; i = nxt[i]) if(i != (e^1)){
leaf = false;
dfs(to[i], i);
merge(hull[x], hull[to[i]]);
}
if(leaf){
hull[x].s.insert({w[e], 2*c[e]});
hull[x].a = c[e], hull[x].b = -c[e] * w[e];
hull[x].num = 2*c[e];
} else{
hull[x].adjust(c[e], 1), hull[x].shift += w[e];
hull[x].b -= hull[x].a * w[e];
}
lb[x] = hull[x].s.begin()->fr, rb[x] = (--hull[x].s.end())->fr;
lb[x] += hull[x].shift, rb[x] += hull[x].shift;
mn[x] = hull[x].a - hull[x].num, mx[x] = hull[x].a;
}
void solution(int x, int e, ll L){
len[e] = w[e];
if(L > rb[x] && mx[x] == c[e]) len[e] += L-rb[x];
else if(L < lb[x] && mn[x] == -c[e]) len[e] -= lb[x]-L;
len[e^1] = len[e];
for(int i = head[x]; ~i; i = nxt[i]) if(i != (e^1))
solution(to[i], i, L-len[e]);
}
int main(){
ios::sync_with_stdio(false);
int T; cin>>T;
while(T--){
cin>>n;
int u, v, _w, _c; init();
rep(i,1,n-1) cin>>u>>v>>_w>>_c, add_e(u, v, _w, _c, 1);
for(int i = head[1]; ~i; i = nxt[i])
dfs(to[i], i), merge(hull[1], hull[to[i]]);
hull[1].adjust(0, 0);
ll L = hull[1].s.empty() ? 0 : (--hull[1].s.end())->fr + hull[1].shift;
cout<< hull[1].a * L + hull[1].b <<endl;
for(int i = head[1]; ~i; i = nxt[i])
solution(to[i], i, L);
rep(i,0,cnt/2) cout<< len[i*2] <<endl;
}
return 0;
}
