[C/C++] 백준 - 9489번 : 사촌

https://www.acmicpc.net/problem/9489

9489번: 사촌

문제

증가하는 정수 수열을 이용해서 트리를 만드는 방법은 다음과 같다.

• 첫 번째 정수는 트리의 루트 노드이다.
• 다음에 등장하는 연속된 수의 집합은 루트의 자식을 나타낸다. 이 집합에 포함되는 수의 첫 번째 수는 항상 루트 노드+1보다 크다.
• 그 다음부터는 모든 연속된 수의 집합은 아직 자식이 없는 노드의 자식이 된다. 그러한 노드가 여러 가지 인 경우에는 가장 작은 수를 가지는 노드의 자식이 된다.
• 집합은 수가 연속하지 않는 곳에서 구분된다.

예를 들어, 수열 1 3 4 5 8 9 15 30 31 32를 위의 규칙을 이용해 트리를 만들면 아래 그림과 같이 된다.

두 노드의 부모는 다르지만, 두 부모가 형제(sibling)일 때 두 노드를 사촌이라고 한다.

수열 특정 노드 번호 k가 주어졌을 때, k의 사촌의 수를 구하는 프로그램을 작성하시오.

 

---

주어진 수열에서 가장 첫 번째 수는 루트노드를 의미하고 그 뒤 숫자는 연속된 수들의 집합으로 다른 형태로 위치가 지정되어있다. 가장 먼저 해야하는 것은 연속된 수들로 집합을 구성하여 문제에서 요구하는대로 트리의 형태를 만들어야한다.

루트노드를 하나 생성해놓고 그 밑으로 트리의 형태가 펼처지는 꼴이다. 그리고 자식 노드가 생성될 때에는 자식이 없는노드 중 가장 작은 수를 가지는 노드의 자식으로 생성된다. 나는 이 부분을 우선순위 큐로 생성이 가능하다고 판단했다. 

 

루트노드를 제외하고 연속적인 부분이 나올 때는 우선순위 큐에서 top에 해당하는 숫자를 가지는 노드의 자식으로 배치를 하고 연속수열이 깨질 경우에는 우선순위 큐에서 새로운 값을 빼서 부모노드로 지정하고 배치해주면된다. 이 부분은 for문과 우선순위 큐로 트리의 형태를 만드는 과정이다.

 

하나의 노드는 하나의 부모만 가지므로 부모노드를 저장하는 것은 배열로 쉽게 구현할 수 있다. 위 과정을 진행하면서 부모노드 번호도 기억해주어야한다.

 

위 부분만 끝나면 이제는 K의 사촌 수만 구하면 된다.

이 부분은 전혀 어렵지 않다. 두 노드의 부모는 다르지만, 두 부모가 형제인 노드만 구하면 되므로 K의 부모노드(이를 parentNode로 칭함) parentNode의 부모노드(grandParentNode)만 구하면 된다. 위에서 만들었던 부모노드를 저장하는 배열로 빠르게 구할 수 있다.

 

이제 grandParentNode를 부모로 가지는 노드 중 parentNode를 제외하고(parentNode가 가지는 자식들은 K와 형제 관계이므로) 나머지 노드가 가지는 자식 수를 모두 더해주면 K의 사촌수가 된다.

#include <algorithm>
#include <iostream>
#include <vector>
#include <queue>
#define fastio ios::sync_with_stdio(false), cin.tie(0), cout.tie(0)
#define ENDL cout << endl
#define endl '\n'

using namespace std;

int N, K, root;
int parent[1000000 + 1];
vector<int> tree[1000000 + 1];

void clear(vector<int> &t) {
	for (auto& ptr : t) {
		tree[ptr].clear();
	}
}

int solution() {
	int upRoot = parent[K];
	int dfsRoot = parent[upRoot];

	int answer = 0;
	for (auto& ptr : tree[dfsRoot]) {
		if (ptr == upRoot) continue;
		answer += tree[ptr].size();
	}

	return answer;
}

void init_Tree(vector<int>& ps) {
	int prev, curParent;
	bool start = false;
	priority_queue<int, vector<int>, greater<int>> pq;
	
	root = ps[0];
	parent[root] = -1;

	pq.push(root);
	for (int i = 1; i < N; i++) {
		if (!start) {
			start = true;
			curParent = pq.top(); pq.pop();
		}
		else {
			// 연속적으로 증가하는 부분 수열이 아닐 경우 재 시작 해야함
			if (ps[i] - prev != 1) {
				curParent = pq.top(); pq.pop();
			}
		}
		prev = ps[i];
		tree[curParent].push_back(ps[i]);
		parent[ps[i]] = curParent;
		pq.push(ps[i]);
	}
}

void input() {
	while (1) {
		cin >> N >> K;
		if (!N && !K) break;

		vector<int> ps(N, 0);
		for (int i = 0; i < N; i++) {
			cin >> ps[i];
		}

		init_Tree(ps);
		cout << solution() << endl;
		clear(ps);
	}
}

int main() {
	fastio;
	input();
	
	return 0;
}