[P2] 백준 15899번 트리와 색깔 C++ 세그먼트 트리, 오일러 경로 테크닉, 머지 소트 트리

리뷰

 

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

서브 트리에서 특정 값 이하의 개수를 구하는 문제

 

 

전역 변수

• N : 정점의 개수를 저장할 변수
• MOD : 모듈러 연산을 위한 변수
• lst : 각 노드의 색을 저장할 배열
• color : 노드 방문 정보를 기준으로 색을 저장할 배열
• edges : 인접 리스트를 저장할 벡터 배열
• n : 정점의 개수
• m : 쿼리의 개수를 저장할 변수
• c : 색의 최대값을 저장할 변수
• it : 트리 탐색 시 진입 시간을 저장할 배열
• ot : 트리 탐색 시 탈출 시간을 저장할 배열
• t : 트리 탐색 시 시간을 저장할 변수
• ans : 정답을 저장할 변수
• tree : 세그먼트 트리 정보를 저장할 배열

 

함수

1. dfs

void dfs(int node, int par)

 

깊이 우선 탐색을 통해 오일러 경로를 최신화 하기 위한 함수

1. 매개 변수로 현대 노드 mode와 부모 노드 par을 전달 받는다.
2. it[node]에 t를 전위 증가한 값을 저장해 준다.
3. color[it[node]]의 값을 현재 노드의 색 lst[node]로 저장해 준다.
4. node의 인접 리스트를 순회하며 재귀를 진행해 준다.
5. 모든 인접 리스트 잡색을 마친 뒤 ot[node]에 t값을 저장해 준다.

 

2. build

void build(int node, int s, int e)

 

세그먼트 트리 정보를 초기화 하기 위한 함수

1. 매개 변수로 노드 정보 node, 탐색 범위 s, e를 전달 받는다.
2. 리프 노드에 도달한 경우 해당 노드에 color[s]혹은 color[e]값을 벡터로 저장해 준다.
3. 아닌 경우 좌, 우 자식 노드로 재귀를 진행해 준다.
4. 재귀를 빠져나온 후 왼쪽 및 오른쪽 자식 노드를 벡터 변수 left, right로 가져온다.
5. 벡터 par을 left, right 크기의 합으로 초기화 해준다.
6. merge함수를 통해 left, right를 병합 정렬하여 par의 begin에 요소를 삽입해 준다.
7. 현재 노드에 par값을 저장해 준다.

 

3. query

int query(int node, int s, int e, int L, int R, int C)

 

서브 트리에서 C이하인 요소 개수를 구하는 함수

1. 매개 변수로 노드 정보 node, 탐색 범위 s, e, 구간 범위 L, R, 값 C를 전달 받는다.
2. 첫 번째 기저 조건으로 범위를 벗어난 경우 0을 리턴해 준다.
3. 두 번째 기저 조건으로 탐색 범위가 구간에 포함될 경우 upper_bound함수를 통해 현재 노드으 벡터에서 C이하인 요소의 개수를 구해 개수를 리턴해 준다.
4. 좌, 우 자식 노드로 재귀 탐색을 진행하고 각각 받은 리턴 값을 변수 left, right에 저장해 준다.
5. left와 right를 더한 값을 리턴해 준다.

 

문제풀이

1. n, m, c값을 입력 받는다.
2. n개의 노드의 색 정보를 lst배열에 입력 받아준다.
3. n - 1개의 간선 정보를 입력 받아 인접 리스트에 양방향으로 추가해 준다.
4. dfs함수에 1, 1을 매개 변수로 전달해 오일러 경로를 최신화 해준다.
5. build 함수를 통해 오일러 경로를 기반으로 세그먼트 트리를 초기화 해준다.
6. m개의 쿼리를 입력 받아 매 쿼리 마다 노드 번호 V, 값 C를 입력 받아준다.
7. query 함수를 통해 it[V] ~ ot[V] 범위에서 C이하인 요소의 개수를 구해 ans에 더해준다.
8. 만약 ans가 MOD보다 커졌을 경우 나머지 연산을 진행해 준다.
9. 모든 탐색을 진행한 후에 ans에 저장된 값을 출력해 준다.

 

트러블 슈팅

1. 어차피 트리이므로 n - 1개의 간선을 단 방향으로 추가했으나 Fail을 받았다.
2. 루트 번호가 1이라는 보장이 없기 때문에 양방향으로 추가하고 dfs내부에 조건을 추가해 주었다.
3. query 함수의 반환값을 벡터로 하고, 매번 merge를 수행해 줬더니 시간 초과를 받았다.
4. 반환값을 int로 변경하고 내부에서 리턴할 때 upper_bound를 수행해 주었더니 AC를 받았다.

 

참고 사항

• 예제는 엣지케이스 탐색에 별 도움이 되지 않는다.

 

정답 코드

#include<iostream>
#include<vector>
#include<algorithm>
#define ll long long
using namespace std;

const int N = 200001;
const int MOD = 1000000007;
int lst[N];
int color[N];
vector<int> edges[N];
int n, m, c;
int it[N];
int ot[N];
int t;
ll ans;
vector<int> tree[N * 4];

void dfs(int node, int par) {
	it[node] = ++t;
	color[it[node]] = lst[node];
	for (int i : edges[node]) {
		if (par == i) continue;
		dfs(i, node);
	}
	ot[node] = t;
}

void build(int node, int s, int e) {
	if (s == e) tree[node] = { color[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] = par;
	}
}

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

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

	cin >> n >> m >> c;
	for (int i = 1; i <= n; ++i) cin >> lst[i];
	for (int i = 0; i < n - 1; ++i) {
		int a, b; cin >> a >> b;
		edges[a].push_back(b);
		edges[b].push_back(a);
	}

	dfs(1, 1);
	build(1, 1, n);

	while (m--) {
		int V, C; cin >> V >> C;
		ans += query(1, 1, n, it[V], ot[V], C);
		if (ans >= MOD) ans %= MOD;
	}
	cout << ans;
}