USACO2022DEC silver1 barn tree 202212月银牌第一题 （DFS）

#include<bits/stdc++.h>

using namespace std;

typedef long long ll;

int n;

const int MAXN=200001;

vector<ll> w(MAXN);

set<int> g[MAXN];

set<pair<int,ll>> gdw[MAXN];

vector<int> indegree(MAXN);

vector<set<int>> d2n(MAXN);

ll sumw,avgw;

vector<tuple<int,int,ll>> record;

int totalmove;

bool visited[MAXN];

ll dfs(int i,int& nodecount){

ll totalw=w[i];

visited[i]=true;

//cout<<"vis:"<<i<<endl;

for(auto x:g[i]){

if(!visited[x]){

//visited[x]=true;

int nc=0;

ll rv=dfs(x,nc);

if(rv>0){

gdw[x].insert(make_pair(i,rv));

indegree[i]++;

}else{

if(rv<0){

gdw[i].insert(make_pair(x,-rv));

indegree[x]++;

}

}

nodecount+=nc;

totalw+=rv+nc*avgw;

}

}

nodecount++;

//cout<<i<<" "<<totalw<<" "<<avgw*nodecount<<" "<<nodecount<<endl;

return totalw-avgw*nodecount;

}

int main(){

cin>>n;

for(int i=1;i<=n;i++){

cin>>w[i];

sumw+=w[i];

}

avgw=sumw/n;

for(int i=1;i<=n-1;i++){

int u,v;

cin>>u>>v;

g[u].insert(v);

g[v].insert(u);

}

int nc=0;

if(dfs(1,nc)!=0){

cout<<"input data error 1\n";

}

if(nc!=n){

cout<<"input data error 2\n";

}

//topology sequence

for(int i=1;i<=n;i++){

d2n[indegree[i]].insert(i);

//cout<<i<<" "<<indegree[i]<<endl;

}

while(d2n[0].size()>0){

int tn=*(d2n[0].begin());

for(auto ip=gdw[tn].begin();ip!=gdw[tn].end();){

auto x=*ip;

int nn=x.first;

ll w=x.second;

d2n[indegree[nn]].erase(nn);

indegree[nn]--;

d2n[indegree[nn]].insert(nn);

totalmove++;

record.push_back(make_tuple(tn,nn,w));

auto tip=ip;

ip++;

gdw[tn].erase(tip);

}

d2n[0].erase(tn);

}

printf("%d\n",totalmove);

for(int i=0;i<record.size();i++){

printf("%d %d %ld\n",get<0>(record[i]),get<1>(record[i]),get<2>(record[i]));

}

    return 0;

}