HackerRank: World codesprint #4 - Roads in HackerLand - study C++ code V

RoadsInHackerLand5A.cpp

#include "bits/stdc++.h"
using namespace std;
#define rep(i,n) for(int (i)=0;(i)<(int)(n);++(i))
#define rer(i,l,u) for(int (i)=(int)(l);(i)<=(int)(u);++(i))
#define reu(i,l,u) for(int (i)=(int)(l);(i)<(int)(u);++(i))
static const int INF = 0x3f3f3f3f; static const long long INFL = 0x3f3f3f3f3f3f3f3fLL;
typedef vector<int> vi; typedef pair<int, int> pii; typedef vector<pair<int, int> > vpii; typedef long long ll;
template<typename T, typename U> static void amin(T &x, U y) { if(y < x) x = y; }
template<typename T, typename U> static void amax(T &x, U y) { if(x < y) x = y; }

struct UnionFind {
	vector<int> data;
	void init(int n) { data.assign(n, -1); }
	bool unionSet(int x, int y) {
		x = root(x); y = root(y);
		if(x != y) {
			if(data[y] < data[x]) swap(x, y);
			data[x] += data[y]; data[y] = x;
		}
		return x != y;
	}
	bool findSet(int x, int y) { return root(x) == root(y); }
	int root(int x) { return data[x] < 0 ? x : data[x] = root(data[x]); }
	int size(int x) { return -data[root(x)]; }
};

vector<int> t_parent;
vi t_ord;

void tree_getorder(const vector<vi> &g, int root) {
	int n = g.size();
	t_parent.assign(n, -1);
	t_ord.clear();

	vector<int> stk; stk.push_back(root);
	while(!stk.empty()) {
		int i = stk.back(); stk.pop_back();
		t_ord.push_back(i);
		for(int j = (int)g[i].size() - 1; j >= 0; j --) {
			int c = g[i][j];
			if(t_parent[c] == -1 && c != root)
				stk.push_back(c);
			else
				t_parent[i] = c;
		}
	}
}

int main() {
	int N; int M;
	while(~scanf("%d%d", &N, &M)) {
		vector<pii> edges(M);
		rep(i, M) {
			int A; int B; int C;
			scanf("%d%d%d", &A, &B, &C), -- A, -- B;
			edges[C] = make_pair(A, B);
		}
		UnionFind uf; uf.init(N);
		vector<vi> g(N);
		map<pii, int> ma;
		rep(i, M) {
			int a, b;
			tie(a, b) = edges[i];
			if(uf.unionSet(a, b)) {
				g[a].push_back(b);
				g[b].push_back(a);
				ma[{a, b}] = ma[{b, a}] = i;
			}
		}
		tree_getorder(g, 0);
		vector<int> subt(N, 1);
		for(int ix = (int)t_ord.size() - 1; ix > 0; -- ix) {
			int i = t_ord[ix], p = t_parent[i];
			subt[p] += subt[i];
		}
		vector<ll> nums(M + 100, 0);
		for(int ix = 1; ix < (int)t_ord.size(); ++ ix) {
			int i = t_ord[ix], p = t_parent[i];
			int ei = ma[{i, p}];
			int x = subt[i], y = N - x;
			nums[ei] += (ll)x * y;
		}
		string ans;
		rep(i, M + 99) {
			nums[i + 1] += nums[i] / 2;
			ans += char('0' + nums[i] % 2);
		}
		while(ans.size() > 1 && ans.back() == '0')
			ans.pop_back();
		reverse(ans.begin(), ans.end());
		puts(ans.c_str());
	}
	return 0;
}