1876E - Ball-Stackable - CodeForces Solution

#include<bits/stdc++.h>
#define fr first
#define sc second
#define pb push_back
using namespace std;const int N=5e5+7;typedef pair<int,int>pa;
int n,f[N],col,i,x,y,z,rt;vector<pa>g[N];vector<int>q;
void dfs1(int x,int fa=0){for(auto y:g[x])if(y.fr!=fa)f[1]+=(y.sc!=-1),dfs1(y.fr, x);}
void dfs2(int x,int fa=0){if(f[x]>f[rt])rt=x;for(auto y:g[x])if(y.fr!=fa)f[y.fr]=f[x]-y.sc,dfs2(y.fr,x);}
void dfs3(int x,int fa=0){
	for(auto y:g[x])if(y.fr!=fa){
		int c;if(y.sc==-1)printf("%d %d %d\n", y.fr, x, c=q.back()),q.pop_back();
		else printf("%d %d %d\n",x,y.fr,++col),q.pb(col);
		dfs3(y.fr,x);if(y.sc==-1)q.push_back(c);else q.pop_back();
	}
}
int main(){
	for(scanf("%d",&n),i=1;i<n;++i)scanf("%d%d%d",&x,&y,&z),g[x].pb(pa(y,z)),g[y].pb(pa(x,-z));
	dfs1(1);dfs2(1);printf("%d\n",f[rt]);dfs3(rt);
}