Indeedなう（本選A） : D - Longest Path

TLE対策に苦戦してしまった。
http://indeednow-finala-open.contest.atcoder.jp/tasks/indeednow_2015_finala_d

問題

N個の頂点と(N-1)個の辺からなるグラフがある。
各辺は有向辺か無向辺である。
また、各辺が無向辺とみなすと、このグラフは木となっている。

N上における２頂点間のパスa→bのうち、経由する頂点の最大値を求めよ。

解法

Editorialでは木DPを使っていたが、自分は別解で解いたのでここではそちらを紹介。

各頂点において、(頂点経由数,直前の頂点)の対を頂点経由数が大きい順に上位2個を管理する。
Priority_queueを使い頂点経由数の多い順に、各頂点から隣接点への移動を行う。

頂点xにおいて、(頂点経由数,直前の頂点)の上位2個を(P1,Y1),(P2,Y2)とする。
xの隣接点yにおいて以下のように隣接点における頂点経由数上位２個を更新する。

• y==Y1ならyに至る最大頂点経由数をP2+1で更新する。
• y!=Y1ならyに至る最大頂点経由数をP2+1で更新する。

この処理で上位２個が入れ替わった場合、頂点yを再探索対象としてPriority_queueに再度入れる。

このようにすると、頂点経由数が多い順に隣接点をどんどんたどり、求める答えが得られる。

int N;
int P[101000],in[101000],mama[101000];
vector<int> E[101000];
set<pair<int,int> > S[101000];
map<int,int> M[101000];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N-1) {
		cin>>x>>y>>j;
		x--,y--;
		E[x].push_back(y);
		if(j==2) E[y].push_back(x);
	}
	
	priority_queue<pair<int,int> > Q;
	FOR(i,N) Q.push(make_pair(0,i));
	
	while(Q.size()) {
		auto e=Q.top();
		int cur=e.second;
		Q.pop();
		if(e.first<mama[cur]) continue;
		
		FORR(tar,E[cur]) {
			x = 0;
			if(S[cur].size()) {
				if(S[cur].rbegin()->second != tar) x=S[cur].rbegin()->first;
				else if(S[cur].size()>1) x=S[cur].begin()->first;
			}
			if(M[tar][cur]<x+1) {
				S[tar].erase(make_pair(M[tar][cur],cur));
				M[tar][cur]=x+1;
				mama[tar]=max(mama[tar],x+1);
				
				if(S[tar].size()<2 || S[tar].begin()->first < x+1) {
					S[tar].insert(make_pair(M[tar][cur],cur));
					if(S[tar].size()>2) S[tar].erase(S[tar].begin());
					Q.push(make_pair(mama[tar],tar));
				}
			}
		}
	}
	cout<<*max_element(mama,mama+N)-1<<endl;
}

まとめ

自分で解いたときは無向辺と有向辺をどう処理するかに迷って、木DPには思いが至らなかった。