AtCoder AGC #004 : D - Teleporter

苦戦はしたけど解けて良かった。
http://agc004.contest.atcoder.jp/tasks/agc004_d

問題

1～NのN個の頂点がある。
初期状態では各頂点iからA[i]に有向辺が振られている。
また、どの頂点からでも、有向辺を何度か辿ると頂点1に戻れることが分かっている。

ここで、任意の頂点から始め、有向辺をちょうどK回たどったとき、頂点1に到達するようにしたい。
A[i]を最少何個修正すれば、そのようなグラフになるか。

解法

まずA[1]=1でなければならない。
頂点1が周期aの閉路中にある場合、そのa個の頂点のうち(a-1)個はK回移動後に頂点1に戻れないため、a=1にしかなり得ないためである。

A[1]=1とすると、問題文の「ちょうどK回」を「K回以内」と言い換えられるようになる。
あとはDFSで葉から順に「今見ている頂点iから最も遠い頂点への距離D[i]」を求めていく。
距離D[i]がK-1に達し、頂点iから頂点1まで1回で到達できないなら、A[i]を1にしなければならない。

int N,K;
int A[101010];
int D[101010];
vector<int> E[101010];
int ret;

int dfs(int cur) {
	if(cur==0) return 0;
	if(D[cur]>=0) return D[cur];
	return D[cur]=dfs(A[cur])+1;
}

int dfs2(int cur) {
	int ma=0;
	FORR(r,E[cur]) ma=max(ma,dfs2(r)+1);
	if(cur==0) return 0;
	
	if(ma==K-1 && A[cur]!=0) ret++, ma=-1;
	return ma;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	FOR(i,N) {
		cin>>A[i];
		A[i]--;
		if(i) E[A[i]].push_back(i);
	}
	
	if(A[0]!=0) ret++;
	dfs2(0);
	cout<<ret<<endl;
}

まとめ

こっち先やればよかった。