백준 10838 - 트리(C++)

문제 링크입니다. https://www.acmicpc.net/problem/10838

10838번: 트리

LCA와 같은 원리로 푸는 문제였습니다.

A부터 B까지 같은 색을 연결하는 것이 문제의 핵심입니다.

 

먼저 A와 A의 부모를 모두 찾아서 True로 저장해줍니다. 다음으로 B와 B의 부모를 찾아보는데 True로 저장된 노드를 살펴봅니다. 만약 True로 저장된 노드가 존재한다면 해당 노드가 A와 B의 공통 부모라는 의미입니다.

문제에서 주어진 조건에 따라 1000개 이하로 무조건 찾을 수 있음으로 반복문을 1000번만 돌려줍니다. 

공통 부모를 찾았으면 A와 B모두 부모까지 전부 같은 색으로 칠해줍니다.

 

노드를 연결하는 방법은 간단합니다.

X번 노드의 부모를 저장하는 배열에 해당하는 값을 넣어주면 됩니다.

 

 A부터 B까지 색이 다른 노드를 찾는 방법은 set을 사용하면 됩니다.

set에 노드들의 색깔을 전부 넣은 뒤 set의 size를 출력해주면 됩니다.

 

 

자세한 것은 코드를 참고해주세요

#define _CRT_SECURE_NO_WARNINGS
#include <iostream>
#include <cstdio>
#include <algorithm>
#include <vector>
#include <queue>
#include <stack>
#include <cmath>
#include <cstring>
#include <string>
#include <set>
#define INF 987654321
#define P pair<int,int>
#define PP pair<int, P>
#define PPP pair<P, P>

using namespace std;

int N, K;
vector<PP> R;
int Node_color[100001];
int Node_parent[100001];

int get_LCA(int x, int y) {
	vector<bool> Find(N, false);
	for (int i = 0; i < 1000; i++) {
		Find[x] = true;
		if (x == 0) break;
		x = Node_parent[x];
	}
	for (int i = 0; i < 1000; i++) {
		if (Find[y]) return y;
		y = Node_parent[y];
	}
	return 0;
}

void solve() {
	for (int i = 0; i < R.size(); i++) {
		int type = R[i].first;
		if (type >= 0) {
			int x = R[i].second.first;
			int y = R[i].second.second;
			int lca = get_LCA(x, y);
			while (x != lca) {
				Node_color[x] = type;
				x = Node_parent[x];
			}
			while (y != lca) {
				Node_color[y] = type;
				y = Node_parent[y];
			}
		}
		else if (type == -1) {
			int x = R[i].second.first;
			int y = R[i].second.second;
			Node_parent[x] = y;
		}
		else {
			set<int> color; // set은 중복이 허용되지 않음으로 사용
			int x = R[i].second.first;
			int y = R[i].second.second;
			int ans = get_LCA(x, y);
			while (x != ans) {
				color.insert(Node_color[x]);
				x = Node_parent[x];
			}
			while (y != ans) {
				color.insert(Node_color[y]);
				y = Node_parent[y];
			}
			cout << color.size() << "\n";
		}
	}
}

int main() {
	cin.tie(0);
	cout.tie(0);
	cin >> N >> K;
	for (int i = 0; i < K; i++) {
		int r, a, b, c;
		cin >> r;
		if (r == 1) {
			cin >> a >> b >> c;
			R.push_back({ c,{a,b} });
		}
		else if (r == 2) {
			cin >> a >> b;
			R.push_back({ -1,{a,b} });
		}
		else {
			cin >> a >> b;
			R.push_back({ -2, {a,b} });
		}
	}
	solve();
	return 0;
}

 

 

질문 및 조언 댓글 남겨주세요