[BOJ 15899 // C++] 트리와 색깔

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

 

이번에 볼 문제는 백준 15899번 문제인 트리와 색깔이다.
문제는 아래 링크를 확인하자.

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

15899번: 트리와 색깔

각 노드에 색이 칠해진 트리와, 한 노드의 특정 숫자 이하인 색이 칠해진 색의 개수를 세는 쿼리가 주어지는 문제이다.

 

한 노드를 루트로 하는 서브트리에 관한 쿼리를 처리해야 하므로 euler tour technique을 이용하여 노드에 index를 부여하자.

 

한편, 모든 쿼리의 총합만을 계산하는 문제이므로 쿼리의 순서를 원하는 대로 바꿔 계산해도 문제가 없다.

색깔을 작은 크기서부터 크기순으로 칠해나가면서 모든 세야하는 색이 다 칠해진 쿼리를 그때그때 처리하여 답에 더해나가자. 색칠된 노드의 개수를 세는 것은 색이 칠해진 노드에 가중치 1을 주어 구간합 세그먼트트리를 이용하여 구현할 수 있다.

 

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

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

struct query {
	int v, c;
	query(int v, int c) {
		this->v = v, this->c = c;
	}
};

bool compc(query q1, query q2) {
	return q1.c < q2.c;
}

vector<query> queries;
vector<int> color[200001];
vector<int> G[200001];
int idx = 1;
int id[200001];
int cnt[200001];

int dfs(int current, int parent) {
	int ret = 0;
	id[current] = idx++;
	for (auto node : G[current]) {
		if (node == parent) continue;
		ret += dfs(node, current);
	}
	cnt[current] = ret;
	return ret + 1;
}

int seg[524289];

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

int rangesum(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 rangesum(L, (L + R) / 2, qL, qR, sI * 2) + rangesum((L + R) / 2 + 1, R, qL, qR, sI * 2 + 1);
}

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

	int N, Q, C; cin >> N >> Q >> C;
	
	for (int i = 1; i <= N; i++) {
		int x; cin >> x;
		color[x].emplace_back(i);
	}

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

	for (int i = 0; i < Q; i++) {
		int v, c; cin >> v >> c;
		queries.emplace_back(query(v, c));
	}

	dfs(1, 1);

	int ans = 0;
	int colorid = 0;
	sort(queries.begin(), queries.end(), compc);
	for (auto &q : queries) {
		int c = q.c;
		while (colorid <= c) {
			for (auto node : color[colorid]) {
				update(1, N, id[node]);
			}
			colorid++;
		}
		ans += rangesum(1, N, id[q.v], id[q.v] + cnt[q.v], 1);
		if (ans > 1000000006) ans -= 1000000007;
	}

	cout << ans;
}