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;
}? ??