TopCoder, SRM 635, Division 2, Level 3 : LonglongestPathTree

問題文

　http://community.topcoder.com/stat?c=problem_statement&pm=13416&rd=16078

問題概要

　 頂点からなる重み付きの木が与えられる。辺に関する情報は要素数が の三つの配列 によって与えられ、 番目の辺は を距離 で結んでいる。
　この木に対し、一つの辺を除去してから、同じ長さの辺を付け加える操作を一回まで行える。このとき、操作後のグラフも木になっていなければならない。
　達成できる最大の直径を求めよ。

解法

　ある辺を除去したとき、グラフは分割されて二つの木になります。辺を繋ぎ直すとき、分割されたそれぞれの木の直径の端点同士を選ぶことで最大の直径を達成できます。辺 を除去したときにできる二つの木の直径をそれぞれ とすれば、繋ぎ直した後の直径は となります。
　さて、木の直径は DFS を二回することで求められる*1ので、一回あたり で計算できます。辺の除去も 通りなので、除去する辺を全通り試して毎回直径を計算しても で問題を解けます。

コード

using LL = long long;

#define REP( i, m, n ) for ( int i = (int)( m ); i < (int)( n ); ++i )
#define ALL( c ) (c).begin(), (c).end()

// 木の直径またはその端点を求める
// DFS 二回
// O( |V| )
class TreeDiameter
{
private:
	const int N;
	std::vector< std::vector< std::pair< int, long long > > > G;
	std::vector< long long > distance;

public:

	TreeDiameter( const int n ) : N( n ), G( N ), distance( N, -1 ) {};

	void connect( const int u, const int v, const int d = 1 )
	{
		G[u].emplace_back( v, d );
		G[v].emplace_back( u, d );

		return;
	}

	long long diameter( const int r = 0 )
	{
		dfs( r, 0, -1 );
		const int u = max_element( distance.begin(), distance.end() ) - distance.begin();
		dfs( u, 0, -1 );
		return *max_element( distance.begin(), distance.end() );
	}

	pair<int,int> endpoints( const int r = 0 )
	{
		dfs( r, 0, -1 );
		const int u = max_element( distance.begin(), distance.end() ) - distance.begin();
		dfs( u, 0, -1 );
		const int v = max_element( distance.begin(), distance.end() ) - distance.begin();
		return make_pair( u, v );
	}


	LL solve2()
	{
		const int v = find( ALL( distance ), -1 ) - distance.begin();
		fill( ALL( distance ), -1 );
		return diameter( v );
	}

private:
	void dfs( const int u, const long long depth, const int p )
	{
		distance[u] = depth;

		for ( const auto &e : G[u] )
		{
			const int v = e.first;
			const long long d = e.second;

			if ( v == p )
			{
				continue;
			}
			dfs( v, depth + d, u );
		}

		return;
	}
};
// TreeDiameter( N )
// connect( a, b, cost )
// diameter() := 直径を返す
// endpoints() := 直径の端点の pair を返す

class LonglongestPathTree
{
public:
	long long getLength( vector <int> A, vector <int> B, vector <int> L )
	{
		const int V = A.size() + 1;

		LL res = 0;
		REP( i, 0, V - 1 )
		{
			TreeDiameter td( V );
			REP( j, 0, V - 1 )
			{
				if ( i == j )
				{
					continue;
				}
				td.connect( A[j], B[j], L[j] );
			}

			LL tmp = L[i] + td.diameter();
			tmp += td.solve2();
			res = max( res, tmp );
		}

		return res;
	}
};

*1:参考：AtCoder Regular Contest #022, C : ロミオとジュリエット