NFLSOJ #12430. 「NOIP2021模拟赛0915长郡」摩拉克斯（主席树）

题面传送门

人傻常数大指我实锤了

首先考虑怎样暴力地求答案。不难发现，如果我们假设放置两个北国银行的位置分别为 \(x,y\)，那么一定存在一条边 \(x’,y’\)，满足 \(x’\) 到 \(x\) 距离更近，\(y’\) 到 \(y\) 距离更近，此时显然不会有居民通过这条边。因此我们可以考虑枚举这条边然后对两个子树分别求解答案，这样我们就成功地将两个北国银行的情况转化为一个北国银行的情况。而一个北国银行的情况是一个类似于带权重心的问题，一个自然的想法是用类似于树的重心的方法求解，不过这样不太好扩展到原问题。我们考察每一条边，设 \(s1,s2\) 为切掉这条边后两个子树的 size，那么可以发现如果 \(s1<s2\) 我们肯定会将北国银行建在 \(s2\) 一侧使居民走的路程更近，反之我们肯定会建在 \(s1\)​ 一侧。因此这条边的贡献就是 \(\min(s1,s2)·w\)，其中 \(w\) 为这条边权值。

这样暴力做是平方的，考虑优化。我们以 \(1\) 为根对整棵树进行一遍 DFS，假设 \(x\) 子树的大小（也就是其中居民个数）为 \(siz_x\)，那么我们断掉 \((x,y)\)（这里假设 \(y\) 是 \(x\) 的父亲）时，\(x\) 一侧的大小就显然是 \(siz_x\)，\(y\) 一侧子树大小是 \(tot-siz_x\)，这样一来 \(x\) 子树内的所有边的贡献就可以写作 \(\sum\limits_{p\in\text{subtree}(x)}\min(siz_p,siz_x-siz_p)·fw_p\)，其中 \(fw_p\) 表示 \(p\) 与其父亲连的边的权值，可以注意到当 \(siz_p\le\dfrac{siz_x}{2}\) 时，\(\min\) 取到左边的 \(siz_p\)，否则 \(\min\) 取到右边的 \(siz_x-siz_p\)，这个可以通过以 \(siz_p\) 为下标建“权值主席树”解决，这样查询时，只用考察 DFS 序在 \([bgt_x,edt_x]\)，且下标 \(\le/>\dfrac{siz_x}{2}\) 中数的个数/和即可计算出，主席树/二维数点/线段树合并解决，下面的代码中使用的主席树所以常数上天。对于 \(y\) 那一侧的边，大部分情况下的贡献与 \(x\) 那一侧的类似，为 \(\min(siz_p,tot-siz_x-siz_p)·fw_p\)，可以用类似的方法计算，唯独 \(y\to 1\) 的链上的贡献与之不同，因此我们考虑采取这样的思想，先假设它的贡献也是 \(\min(siz_p,tot-siz_x-siz_p)·fw_p\)，然后再一遍 DFS 把它的贡献扣掉并加上真正的贡献，DFS 时计算它原来的贡献可以维护两个 BIT 并在回溯时把值撤销掉，加上它原来的贡献的部分可以通过开一个栈维护其祖先序列，在祖先序列上二分端点+前缀和解决（貌似官方题解给的双针复杂度是错的？一个扫帚即可卡掉。）。时空复杂度均为 \(n\log n\)。

const int MAXN=1e5;
const int MAXP=MAXN*22;
const ll INF=9e18;
int n,a[MAXN+5],hd[MAXN+5],to[MAXN*2+5],nxt[MAXN*2+5],val[MAXN*2+5],ec=0;
void adde(int u,int v,int w){to[++ec]=v;val[ec]=w;nxt[ec]=hd[u];hd[u]=ec;}
ll siz[MAXN+5],tot=0;
int bgt[MAXN+5],edt[MAXN+5],tim=0,edw[MAXN+5],rid[MAXN+5];
ll key[MAXN+5],uni[MAXN+5];int num=0,p[MAXN+5];
void discrete(){
	for(int i=1;i<=n;i++) key[i]=siz[i];
	sort(key+1,key+n+1);key[0]=-1;
	for(int i=1;i<=n;i++) if(key[i]^key[i-1]) uni[++num]=key[i];
}
int getpos(ll v){return lower_bound(uni+1,uni+num+1,v)-uni;}
void dfs0(int x,int f){
	siz[x]=a[x];rid[bgt[x]=++tim]=x;
	for(int e=hd[x];e;e=nxt[e]){
		int y=to[e],z=val[e];if(y==f) continue;
		dfs0(y,x);siz[x]+=siz[y];edw[y]=z;
	} edt[x]=tim;
}
struct chairman_tree{
	struct node{int ch[2];ll sum;} s[MAXP+5];
	int rt[MAXN+5],ncnt=0;
	void reset(){
		for(int i=0;i<=n;i++) rt[i]=0;
		for(int i=1;i<=ncnt;i++) s[i].ch[0]=s[i].ch[1]=s[i].sum=0;
		ncnt=0;
	}
	int& operator [](int x){return rt[x];}
	int modify(int k,int l,int r,int p,ll v){
		int z=++ncnt;s[z]=s[k];s[z].sum+=v;
		if(l==r) return z;int mid=l+r>>1;
		if(p<=mid) s[z].ch[0]=modify(s[k].ch[0],l,mid,p,v);
		else s[z].ch[1]=modify(s[k].ch[1],mid+1,r,p,v);
		return z;
	}
	ll query(int k,int l,int r,int ql,int qr){
		if(!k||ql>qr) return 0;
		if(ql<=l&&r<=qr) return s[k].sum;
		int mid=l+r>>1;
		if(qr<=mid) return query(s[k].ch[0],l,mid,ql,qr);
		else if(ql>mid) return query(s[k].ch[1],mid+1,r,ql,qr);
		else return query(s[k].ch[0],l,mid,ql,mid)+query(s[k].ch[1],mid+1,r,mid+1,qr);
	}
} w,sw;
struct fenwick{
	ll t[MAXN+5];
	void reset(){memset(t,0,sizeof(t));}
	void add(int x,ll v){for(int i=x;i<=num;i+=(i&(-i))) t[i]+=v;}
	ll query(int x){ll ret=0;for(int i=x;i;i&=(i-1)) ret+=t[i];return ret;}
	ll query_range(int l,int r){return query(r)-query(l-1);}
} fw,fsw;
ll res[MAXN+5];
ll Ssw[MAXN+5],Sw[MAXN+5];
ll calc(int l,int r,int ps,ll sz){
	if(l>r) return 0;
	ll sum=sw.query(sw[r],1,num,1,ps-1)-((l-1==0)?0:sw.query(sw[l-1],1,num,1,ps-1));
	sum+=sz*(w.query(w[r],1,num,ps,num)-((l-1==0)?0:w.query(w[l-1],1,num,ps,num)));
	sum-=sw.query(sw[r],1,num,ps,num)-((l-1==0)?0:sw.query(sw[l-1],1,num,ps,num));
	return sum;
}
int stk[MAXN+5],tp=0;
ll sm[MAXN+5],smw[MAXN+5];
void dfs(int x,int f){
	stk[++tp]=x;
	if(x^1){
		ll out=tot-siz[x];
		int ps=upper_bound(uni+1,uni+num+1,out>>1)-uni;
		res[x]-=fsw.query_range(1,ps-1);
		res[x]-=out*fw.query_range(ps,num);
		res[x]+=fsw.query_range(ps,num);
		int v=p[x];
		fw.add(v,edw[x]);fsw.add(v,edw[x]*siz[x]);
		int l=2,r=tp-1,p=1;
		while(l<=r){
			int mid=l+r>>1;
			if(tot-siz[stk[mid]]<=(out>>1)) p=mid,l=mid+1;
			else r=mid-1;
		}
		res[x]+=sm[stk[p]]*tot-smw[stk[p]];
		res[x]+=smw[stk[tp-1]]-smw[stk[p]]-siz[x]*(sm[stk[tp-1]]-sm[stk[p]]);
	}
	for(int e=hd[x];e;e=nxt[e]){
		int y=to[e];if(y==f) continue;
		sm[y]=sm[x]+edw[y];smw[y]=smw[x]+edw[y]*siz[y];
		dfs(y,x);
	} --tp;
	if(x^1){
		int v=p[x];
		fw.add(v,-edw[x]);fsw.add(v,-edw[x]*siz[x]);
	}
}
void clear(){
	for(int i=1;i<=n;i++) sm[i]=Ssw[i]=Sw[i]=res[i]=hd[i]=siz[i]=edw[i]=bgt[i]=edt[i]=key[i]=uni[i]=0;
	tp=tim=num=ec=tot=0;w.reset();fw.reset();sw.reset();fsw.reset(); 
}
int Casen=0;
void solve(){
	scanf("%d",&n);clear();
	for(int i=1;i<=n;i++) scanf("%d",&a[i]),tot+=a[i];
	for(int i=1,u,v,w;i<n;i++){scanf("%d%d%d",&u,&v,&w);adde(u,v,w);adde(v,u,w);}
	dfs0(1,0);discrete();
	for(int i=1;i<=n;i++) p[i]=getpos(siz[i]);
	for(int i=1;i<=n;i++){
		int x=rid[i],v=p[x];
		w[i]=w.modify(w[i-1],1,num,v,edw[x]);
		sw[i]=sw.modify(sw[i-1],1,num,v,edw[x]*siz[x]);
	}
	for(int i=2;i<=n;i++) Sw[p[i]]+=edw[i],Ssw[p[i]]+=edw[i]*siz[i];
	for(int i=1;i<=num;i++) Sw[i]+=Sw[i-1],Ssw[i]+=Ssw[i-1];
	for(int i=2;i<=n;i++){
		ll sz=siz[i],out=tot-sz;
		int ps1=upper_bound(uni+1,uni+num+1,sz>>1)-uni;
		int ps2=upper_bound(uni+1,uni+num+1,out>>1)-uni;
		res[i]=calc(bgt[i],edt[i],ps1,sz);
		res[i]+=Ssw[ps2-1]+out*(Sw[num]-Sw[ps2-1])-Ssw[num]+Ssw[ps2-1];
		res[i]-=calc(bgt[i],edt[i],ps2,out);
	}
	dfs(1,0);ll ans=INF;
	for(int i=2;i<=n;i++) chkmin(ans,res[i]);
	printf("%lld\n",ans);
}
int main(){
	freopen("morax.in","r",stdin);
	freopen("morax.out","w",stdout);
	int qu;scanf("%d",&qu);
	while(qu--) solve();
	return 0;
}
/*
1
9
5 1 9 8 5 6 3 5 8
1 2 5
2 3 6
3 4 9
4 5 1
5 6 8
6 7 9
7 8 1
8 9 10
*/