백준15899 트리와 색깔

작성일 | In

문제 링크

  • http://icpc.me/15899

문제 출처

  • 2018 UCPC 예선 F번

사용 알고리즘

  • merge sort tree
  • dfs

풀이

DFS Ordering을 매겨서 in[v]와 out[v]를 미리 구해주면, 서브트리에 대한 쿼리를 선형 구조에서 구간에 대한 쿼리로 바꿀 수 있습니다.
dfs ordering을 매긴 다음에 이 문제 처럼 잘 해주면 쉽게 풀 수 있습니다.

전체 코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
#include <bits/stdc++.h>
using namespace std;

const int mod = 1e9+7;
int n, m, c;
int color[202020];
vector<int> g[202020];

int in[202020];
int out[202020];
int arr[202020];

int pv = 0;

void dfs(int now){
	in[now] = ++pv;
	arr[pv] = color[now];
	for(auto nxt : g[now]){
		if(in[nxt]) continue;
		dfs(nxt);
	}
	out[now] = pv;
}

vector<int> tree[808080];

void update(int node, int s, int e, int x){
	if(e < x || x < s) return;
	tree[node].push_back(arr[x]);
	if(s ^ e){
		int mid = s + e >> 1;
		update(node*2, s, mid, x);
		update(node*2+1, mid+1, e, x);
	}
}

int get(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].size() - (tree[node].end() - upper_bound(tree[node].begin(), tree[node].end(), x));
	int mid = s + e >> 1;
	return get(node*2, s, mid, l, r, x) + get(node*2+1, mid+1, e, l, r, x);
}

int main(){
	ios_base::sync_with_stdio(0); cin.tie(0);
	cin >> n >> m >> c;
	for(int i=1; i<=n; i++) cin >> color[i];
	for(int i=1; i<n; i++){
		int a, b; cin >> a >> b;
		g[a].push_back(b);
		g[b].push_back(a);
	}
	dfs(1);
	for(int i=1; i<=n; i++){
		update(1, 1, n, i);
	}
	for(int i=0; i<808080; i++) sort(tree[i].begin(), tree[i].end());
	int ans = 0;
	while(m--){
		int a, b; cin >> a >> b;
		int now = get(1, 1, n,in[a], out[a], b);
		ans += now;
		ans %= mod;
	}
	cout << ans;
	return 0;
}