[BOJ 15899] 트리와 색깔(c++)

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

---

문제 개요

N개의 정점과 N-1개의 간선으로 이루어진 트리에서 각 정점은 1 이상 C 이하의 색깔을 가지고 있습니다. M개의 질의 f(v,c)가 주어질 때, 각 질의에 대한 답을 계산하는 문제입니다.

• f(v, c): 정점 v를 루트로 하는 서브트리에서 색깔이 c 이하인 정점의 개수

(1 ≤ N, M ≤ 2×10⁵)

(1 ≤ C ≤ N)

 

 

 

 

---

문제 풀이

이 문제는 오일러 투어 테크닉으로 서브트리를 구간으로 변환한 후 머지소트 트리를 통해 문제를 해결하였습니다.

 

DFS를 통해 각 정점의 서브트리를 일차원 배열의 구간으로 표현하였습니다. 정점 v의 서브트리는 arr[idx[v].first ~ idx[v].second]입니다.

int i;    // 변환될 인덱스를 저장할 정수형 i
int colors[200001];   // 입력으로 받은 색깔 정보 배열
int arr[200001];	  // 오일러 투어 테크닉으로 변환된 인덱스에 기존 색깔 정보를 담을 배열
pair<int, int> idx[200001];  // first는 서브트리의 루트, second는 서브 트리의 마지막 노드

void dfs(int x, int pre) {
	arr[i] = colors[x];
	idx[x].first = i++;
	for (int y : v[x]) {
		if (y == pre) continue;
		dfs(y, x);
	}
	idx[x].second = i - 1;
}

머지 소트 트리로 구간 쿼리를 처리하였습니다. 세그먼트 트리의 각 노드에 해당 구간의 색깔 값을 정렬하여 저장하였고, upper_bound를 통해 c 이하인 원소의 개수를 계산하였습니다. 

 

시간복잡도는 머지 소트 트리를 구성하고, 쿼리당 세그먼트 트리 탐색 + 이분탐색이 이뤄지므로 O(N log N + M log² N)입니다.

---

코드

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

const int mx = 1000000007;

int n, m, k, i;
long long res;
int colors[200001];
int arr[200001];
pair<int, int> idx[200001];

vector<int> v[200001];
vector<int> seg[200001 * 4];

void dfs(int x, int pre) {
	arr[i] = colors[x];
	idx[x].first = i++;
	for (int y : v[x]) {
		if (y == pre) continue;
		dfs(y, x);
	}
	idx[x].second = i - 1;
}

void mergeSortInit(int x, int s, int e) {
	if (s == e) {
		seg[x].push_back(arr[s]);
		return;
	}
	int mid = s + (e - s) / 2;
	seg[x].resize(e - s + 1);
	mergeSortInit(x * 2, s, mid);
	mergeSortInit(x * 2 + 1, mid + 1, e);
	merge(seg[x * 2].begin(), seg[x * 2].end(), seg[x * 2 + 1].begin(), seg[x * 2 + 1].end(), seg[x].begin());
}

int query(int x, int s, int e, int l, int r, int v) {
	if (r < s || e < l) return 0;
	if (l <= s && e <= r) {
		return upper_bound(seg[x].begin(), seg[x].end(), v) - seg[x].begin();
	}
	int mid = s + (e - s) / 2;
	return query(x * 2, s, mid, l, r, v) + query(x * 2 + 1, mid + 1, e, l, r, v);
}

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);
	cout.tie(0);
	cin >> n >> m >> k;
	for (int i = 1; i <= n; i++) {
		cin >> colors[i];
	}
	for (int i = 1; i < n; i++) {
		int x, y;
		cin >> x >> y;
		v[x].push_back(y);
		v[y].push_back(x);
	}
	dfs(1, 0);
	mergeSortInit(1, 0, n - 1);
	for (int i = 0; i < m; i++) {
		int x, y;
		cin >> x >> y;
		int ret = query(1, 0, n - 1, idx[x].first, idx[x].second, y) % mx;
		res += ret;
		res %= mx;
	}
	cout << res << '\n';
	return 0;
}