POJ 1741: Tree

http://poj.org/problem?id=1741

重心分解。解法は蟻本にある。

適当に整備して使いやすいようにした。updateだけいじればよくなっている。

updateも重心を根とする木について考えれば良いので楽。

namespace CD { //centroid decomposition

	struct edge { int to, length; };

	vector<edge> G[MAX_N];
	bool centroid[MAX_N];
	int subtree_size[MAX_N];

	void init(int n) {
		rep(i, 0, n) G[i].clear();
	}
	void add_edge(int a, int b, int c) {
		G[a].push_back((edge){b, c});
		G[b].push_back((edge){a, c});
	}

	int compute_subtree_size(int v, int p) {
		int c = 1;
		rep(i, 0, (int)G[v].size()) {
			int w = G[v][i].to;
			if(w == p || centroid[w]) continue;
			c += compute_subtree_size(G[v][i].to, v);
		}
		subtree_size[v] = c;
		return c;
	}

	//(maxnum of subtree vertexs, centroid vertex)
	pi search_centroid(int v, int p, int t) { 
		pi res = pi(inf, -1);
		int s = 1, m = 0;
		rep(i, 0, (int)G[v].size()) {
			int w = G[v][i].to;
			if(w == p || centroid[w]) continue;

			MIN(res, search_centroid(w, v, t));

			MAX(m, subtree_size[w]);
			s += subtree_size[w];
		}
		MAX(m, t - s);
		MIN(res, pi(m, v));
		return res;
	}

	void update(int s);

	void solve_subproblem(int v) {
		compute_subtree_size(v, -1);
		int s = search_centroid(v, -1, subtree_size[v]).sec;
		centroid[s] = true;
		//(1)
		rep(i, 0, (int)G[s].size()) {
			if(centroid[G[s][i].to]) continue;
			solve_subproblem(G[s][i].to);
		}
		update(s);
		centroid[s] = false;
	}

	////////////////////////////////////////////// edit from here
	
	//distance from centroid
	void enumerate_pathes(int v, int p, int d, vi &ds) {
		ds.push_back(d);
		rep(i, 0, (int)G[v].size()) {
			int w = G[v][i].to;
			if(w == p || centroid[w]) continue;
			enumerate_pathes(w, v, d + G[v][i].length, ds);
		}
	}

	//num of pairs whose lengths is less than k
	int count_pairs(vi &ds) {
		int res = 0;
		sort(ds.begin(), ds.end());
		int j = ds.size();
		rep(i, 0, (int)ds.size()) {
			while(j > 0 && ds[i] + ds[j - 1] > K) j--;
			res += j - (j > i ? 1 : 0);
		}
		return res / 2;
	}

	void update(int s) {
		vi ds;
		ds.push_back(0);
		rep(i, 0, (int)G[s].size()) {
			if(centroid[G[s][i].to]) continue;

			vector<int> tds;
			enumerate_pathes(G[s][i].to, s, G[s][i].length, tds);

			ans -= count_pairs(tds);
			ds.insert(ds.end(), tds.begin(), tds.end());
		}
		ans += count_pairs(ds);
	}
}

蟻本もやっと4章。4章は純愛。