[P2] 백준 14446번 Promotion Counting C++ 머지 소트 트리, 세그먼트 트리, 오일러 경로 테크닉

리뷰

 

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

간만에 풀어본 머지 소트 트리 문제였다.

 

 

전역 변수

• N : 배열의 최대 크기를 정의할 상수 변수
• v : 각 노드의 가중치를 저장할 배열
• tree : 머지 소트 트리 정보를 저장할 배열
• it : 각 노드의 입장 시간을 저장할 배열
• ot : 각 노드의 퇴장 시간을 저장할 배열
• arr : 시간별 노드의 가중치를 저장할 배열
• t : 탐색 시간을 저장할 변수
• n : 노드의 개수를 저장할 변수
• edges : 노드의 인접리스트를 저장할 배열

 

함수

1. dfs

void dfs(int cn) {
	it[cn] = ++t;
	arr[t] = v[cn];
	for (int nn : edges[cn]) dfs(nn);
	ot[cn] = t;
}

 

오일러 경로를 통해 노드의 진입 it, ot, arr배열을 초기화 하기 위한 함수

 

2. build

void build(int node, int s, int e) {
	if (s == e) tree[node] = { arr[s] };
	else {
		int mid = (s + e) / 2;
		build(node * 2, s, mid);
		build(node * 2 + 1, mid + 1, e);
		vector<int>& left = tree[node * 2];
		vector<int>& right = tree[node * 2 + 1];
		vector<int> par(left.size() + right.size());
		merge(left.begin(), left.end(), right.begin(), right.end(), par.begin());
		tree[node] = move(par);
	}
}

 

머지 소트 트리의 초기화를 위한 함수

 

3. query

int query(int node, int s, int e, int L, int R, int x) {
	if (R < s || e < L) return 0;
	if (L <= s && e <= R) return tree[node].end() - upper_bound(tree[node].begin(), tree[node].end(), x);
	int mid = (s + e) / 2;
	return query(node * 2, s, mid, L, R, x) + query(node * 2 + 1, mid + 1, e, L, R, x);
}

 

현재 노드의 진입/퇴장 구간에서 가중치보다 큰 노드의 개수를 구하기 위한 함수

 

 

문제풀이

1. n값을 입력 받고, n개의 노드 가중치를 입력 받아 v배열을 초기화한다.
2. n-1개의 간선 정보를 입력 받아 인접 리스트 edges를 초기화한다.
3. dfs함수를 통해 오일러 경로를 탐색하며 배열 it, ot, arr을 초기화한다.
4. build함수를 통해 머지 소트 트리 초기화를 진행한다.
5. 1번 노드부터 n번 노드까지 순회하며 자신의 진입/퇴장 구간과 가중치를 query함수의 매개변수로 전달해 호출한다.
6. query함수 내부에서 upper_bound함수를 통해 자신보다 큰 노드의 개수를 구해 리턴한다.
7. query함수의 리턴값을 출력하고 개행 문자를 삽입해 줄 바꿈을 삽입한다.

 

트러블 슈팅

1. 오일러 경로를 구할때 별도의 배열에 시간 별 가중치에 해당하는 역 구조 매핑을 해주지 않아 WA를 받았다.
2. 배열 arr에 시간 별 가중치를 매핑하고, build함수 내부에서 리프 노드에 활용해 AC를 받았다.

 

참고 사항

1. query에서 upper_bound를 사용하지 않고, 계속 upper_bound를 사용해줄 경우 시간이 터질 것 같다.
2. build함수 내부에서도 각 벡터 정보를 레퍼런스로 받고, merge 후 move를 통해 복사 없이 바로 값을 이동해주었다.

 

정답 코드

#include<iostream>
#include<algorithm>
#include<vector>
using namespace std;

const int N = 1e5 + 1;
int v[N];
vector<int> tree[N * 4];
int it[N];
int ot[N];
int arr[N];
int t;
int n;
vector<int> edges[N];

void dfs(int cn) {
	it[cn] = ++t;
	arr[t] = v[cn];
	for (int nn : edges[cn]) dfs(nn);
	ot[cn] = t;
}

void build(int node, int s, int e) {
	if (s == e) tree[node] = { arr[s] };
	else {
		int mid = (s + e) / 2;
		build(node * 2, s, mid);
		build(node * 2 + 1, mid + 1, e);
		vector<int>& left = tree[node * 2];
		vector<int>& right = tree[node * 2 + 1];
		vector<int> par(left.size() + right.size());
		merge(left.begin(), left.end(), right.begin(), right.end(), par.begin());
		tree[node] = move(par);
	}
}

int query(int node, int s, int e, int L, int R, int x) {
	if (R < s || e < L) return 0;
	if (L <= s && e <= R) return tree[node].end() - upper_bound(tree[node].begin(), tree[node].end(), x);
	int mid = (s + e) / 2;
	return query(node * 2, s, mid, L, R, x) + query(node * 2 + 1, mid + 1, e, L, R, x);
}

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);
	cout.tie(0);

	cin >> n;
	for (int i = 1; i <= n; ++i) cin >> v[i];
	for (int i = 2; i <= n; ++i) {
		int p; cin >> p;
		edges[p].push_back(i);
	}
	dfs(1);
	build(1, 1, n);

	for (int i = 1; i <= n; ++i) {
		cout << query(1, 1, n, it[i], ot[i], v[i]) << "\n";
	}
}