취미로 코딩하기

※ 글쓴이는 취미로 코딩을 익혀보는 사람이라 정확하지 않은 내용을 담고 있을 수 있다 ※

이번에 볼 문제는 백준 14446번 문제인 Promotion Counting이다.
문제는 아래 링크를 확인하자.

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

주어지는 rooted tree의 각 노드에 대하여 자손들 중 해당 노드의 가중치보다 높은 가중치를 갖는 노드의 개수를 구하는 문제이다.

주어지는 트리의 모든 자손 노드를 대상으로 하는 쿼리는 Euler Tour Technique을 이용해 노드의 번호를 재설정하고 각 노드의 자손의 수를 전처리한 다음 세그먼트트리를 이용해 쉽게 해결할 수 있다.

노드의 가중치가 큰 노드부터 차례대로 살펴보며 기존에 등장한 노드를 자손으로 얼마나 포함하고 있는지 계산 및 해당 노드를 세그먼트트리에 업데이트하는 것을 반복해 문제를 해결하자.

이를 구현할 때, 같은 가중치를 갖는 노드에 주의하자. Euler Tour Technique을 통해 새로 부여된 번호를 기준으로 큰 노드는 작은 노드의 조상이 절대 될 수 없음을 이용하면 특별한 예외처리 없이 정렬만 해도 문제를 충분히 해결할 수 있다.

아래는 제출한 소스코드이다.

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

int N;
int seg[262145];

void upd(int qI) {
	int L = 1, R = N, sI = 1;
	while (L < R) {
		seg[sI]++;
		int mid = (L + R) / 2;
		if (mid < qI) L = mid + 1, sI = sI * 2 + 1;
		else R = mid, sI = sI * 2;
	}
	seg[sI]++;
}

int qry(int L, int R, int qL, int qR, int sI) {
	if (R < qL || qR < L) return 0;
	if (qL <= L && R <= qR) return seg[sI];
	return qry(L, (L + R) / 2, qL, qR, sI * 2) + qry((L + R) / 2 + 1, R, qL, qR, sI * 2 + 1);
}

vector<int> G[100001];
int idx, dfi[100001], cnt[100001], inv[100001];
int dfs(int cur, int par) {
	int ret = 0;
	idx++;
	dfi[cur] = idx, inv[idx] = cur;

	for (auto& nxt : G[cur]) {
		if (nxt == par) continue;
		ret += dfs(nxt, cur);
	}
	cnt[dfi[cur]] = ret;
	return ret + 1;
}

pair<int, int> arr[100000];
int ans[100001];

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

	cin >> N;
	for (int i = 0; i < N; i++) {
		int x; cin >> x;
		arr[i] = make_pair(-x, i + 1);
	}

	for (int i = 2; i <= N; i++) {
		int x; cin >> x;
		G[i].emplace_back(x);
		G[x].emplace_back(i);
	}
	dfs(1, 1);

	for (int i = 0; i < N; i++) {
		arr[i].second = dfi[arr[i].second];
	}
	sort(arr, arr + N);

	for (int i = 0; i < N; i++) {
		int ii = arr[i].second;
		ans[inv[ii]] = qry(1, N, ii, ii + cnt[ii], 1);
		upd(ii);
	}

	for (int i = 1; i <= N; i++) cout << ans[i] << '\n';
}