FHC 2018 Round1 B - Ethan Traverses a Tree

問題

Ethanは「各頂点に数字が割り振られている木を与えられ、通る頂点の数字をpre-order travesal(通りがけ順)に並べた数列を作る」という課題を与えられた
しかし間違えてpost-order traversal(帰りがけ順)の数列を作るアルゴリズムを作ってしまった

1) Ethanのアルゴリズムでも正解になる
2) 頂点に割り振る数字は[1,k]のどれかで、k個の数字全てを1度は使う
これらの条件を満たすように頂点に数字を割り振ることは可能か判定

解法

まず頂点番号をそのまま割り振り、実際にpre-orderとpost-order両方の数列(それぞれx,yとする)を作る
条件1を満たすには[0,n)の任意のiについてx[i] = y[i]である必要がある
これはunion findでx[i]とy[i]を全てuniteしておけば、同じ値を書かないといけない頂点のグループが全てわかる
また、条件2はこのグループの個数がk以上であれば満たせる
(i個目のグループに属している頂点にはmin(i,k)を書けばいい)

実装

・unionfindで各頂点の親を探し、その値をsetに格納してグループの個数を調べる
・setに入れた各値を一個一個見ていきそのグループに書く数字rを決め、mapに入れておく

struct UF {
	vector<int> p;
	UF(int n) : p(n, -1) {};
	bool unite(int u, int v) {
		if ((u = find(u)) == (v = find(v))) return false;
		if (p[u] > p[v]) swap(u, v);
		p[u] += p[v]; p[v] = u;
		return true;
	}
	bool same(int u, int v) { return find(u) == find(v); }
	int find(int u) { return p[u] < 0 ? u : p[u] = find(p[u]); }
	int usize(int u) { return -p[find(u)]; }
};

int T, n, k, a, b;
vector<int> e[2020], x, y;

void pre(int v) {
	x.push_back(v);
	each(i,e[v]) pre(i);
}

void pos(int v) {
	each(i,e[v]) pos(i);
	y.push_back(v);
}

signed main() {
	cin >> T;
	reps(TT,1,T+1) {
		cout << "Case #" << TT << ": ";
		cin >> n >> k;
		rep(i,n) e[i].clear();
		rep(i,n) {
			cin >> a >> b;
			if (a) e[i].push_back(a-1);
			if (b) e[i].push_back(b-1);
		}
		
		x.clear(), y.clear();
		pre(0), pos(0);
		UF uf(n);
		rep(i,n) uf.unite(x[i],y[i]);
		set<int> ps;
		rep(i,n) ps.insert(uf.find(i));
		
		if (ps.size()<k) {
			cout << "Impossible" << endl;
			continue;
		}
		
		int r = 1;
		map<int,int> m;
		each(i,ps) {
			m[i] = r;
			if (r<k) r++;
		}
		rep(i,n) cout << m[uf.find(i)] << (i==n-1 ? '\n' : ' ');
	}
}