백준 15899 c++ "트리와 색깔" -PlusUltraCode-

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

 

[필자 사고]

이 문제는 세그먼트 자료구조를 이용하여 푸는 문제이다.

먼저 위계질서가 정해져 있으므로 DFS를 이용하여 index를 세그먼트화 해야 된다.

또한 자신의 값들도 있으므로 세그먼트화 된 index에 맞게 해당 값들도 저장해야 된다.

 

C보다 적은값의 갯수를 묻고 있으므로 c+1로 lowerBound하면 갯수를 구할 수 있다.

 

또한 정점 V의 서브트리라고 했으니 정점 V를 포함하여 index[a].first 부터 시작하면 된다.

 

정밀하지 않으면 실수가 많은 문제이다. 또한 많은 알고리즘이 들어가 있어서 배울수 있는 문제이다.

[코드 해설]

2. 입력 처리 (Input 함수)

Input() 함수는 사용자 입력을 받아 데이터를 초기화합니다.

• N, M, C 값을 입력받습니다.
• arr 배열에 각 노드의 값을 저장합니다.
• 이후 간선 정보를 입력받아 인접 리스트 list를 생성합니다.
• list는 각 노드의 연결 정보를 저장하며, 무방향 그래프로 구성됩니다.

---

3. DFS를 이용한 인덱스 계산 (Dfs 함수)

Dfs(num, parent) 함수는 DFS를 수행하며 각 노드의 방문 순서를 기록합니다.

• index[num].first: num 노드의 DFS 방문 시작 시점.
• index[num].second: num 노드의 DFS 방문 종료 시점.
• arr2[cnt]: cnt번째 방문한 노드의 값을 저장.

DFS를 통해 트리를 선형 구조로 변환하여 세그먼트 트리의 입력으로 사용합니다.

---

4. 세그먼트 트리 초기화 (Init_Tree 함수)

Init_Tree(node, start, end) 함수는 세그먼트 트리를 초기화합니다.

• 리프 노드(start == end)의 경우, arr2[start] 값을 그대로 저장합니다.
• 내부 노드는 두 자식 노드의 값들을 병합하여 정렬된 배열로 저장합니다.
• 병합은 merge 함수를 사용하며, 병합된 배열은 세그먼트 트리의 각 노드에 저장됩니다.

---

5. 쿼리 처리 (Query 함수)

Query(node, start, end, left, right, val) 함수는 주어진 범위에서 특정 조건을 만족하는 값의 개수를 구합니다.

• 쿼리 범위가 노드 범위와 겹치지 않으면 0을 반환합니다.
• 쿼리 범위가 노드 범위에 완전히 포함되면, tree[node]에서 val보다 작은 값의 개수를 이진 탐색으로 계산합니다.
  • lower_bound를 사용하여 효율적으로 개수를 구합니다.
• 범위가 겹칠 경우, 자식 노드로 나누어 재귀적으로 호출하며 결과를 합산합니다.

---

6. 게임 실행 (Game_Start 함수)

Game_Start() 함수는 전체 로직을 실행합니다.

• DFS로 트리를 선형 구조로 변환합니다.
• 세그먼트 트리를 초기화합니다.
• M개의 쿼리를 처리하여 결과를 계산합니다.
  • 각 쿼리에서 특정 노드의 서브트리 내 값 중 기준 b를 넘지 않는 값의 개수를 계산합니다.
• 모든 쿼리 결과를 합산하여 출력합니다.

---

7. 메인 함수

main() 함수는 프로그램의 시작점으로, 다음 작업을 수행합니다.

• 입력을 처리하기 위해 Input() 호출.
• 게임 로직을 수행하기 위해 Game_Start() 호출.
• 최종 결과를 출력합니다.

[소스 코드]

#include <iostream>
#include <vector>
#include <algorithm>
#define ll long long
#define all(v) v.begin(), v.end()

using namespace std;

int N, M, C;
vector<vector<ll>> tree;
vector<vector<int>> list;
vector<pair<int, int>> index;
vector<int> arr;
vector<int> arr2;
int cnt = 0;

void Input() {
	cin >> N >> M >> C;
	tree.resize(4 * N);
	list.resize(N + 1);
	index.resize(N + 1);
	arr.resize(N + 1);
	arr2.resize(N + 1);

	for (int i = 1; i <= N; i++) {
		cin >> arr[i];
	}
	for (int i = 1; i <N; i++) {
		int s, e;
		cin >> s >> e;
		list[s].push_back(e);
		list[e].push_back(s);
	}
}
void Dfs(int num, int parent) {
	index[num].first = ++cnt;
	arr2[cnt] = arr[num];
	for (int next : list[num]) {
		if (next != parent) {
			Dfs(next,num);
		}
	}

	index[num].second = cnt;
}
void Init_Tree(int node, int start, int end) {
	if (start == end) {
		tree[node].push_back(arr2[start]);
		return;
	}
	tree[node].resize(end - start + 1);
	int mid = (start + end) / 2;
	Init_Tree(node * 2, start, mid);
	Init_Tree(node * 2 + 1, mid + 1, end);

	merge(all(tree[node*2]), all(tree[node*2+1]), tree[node].begin());
}
ll Query(int node, int start, int end, int left, int right, ll val) {
	if (end < left || right < start)return 0;

	if (left <= start && end <= right) {
		return lower_bound(all(tree[node]), val + 1) - tree[node].begin();
	}

	int mid = (start + end) / 2;
	return Query(node * 2, start, mid, left, right, val) +
		Query(node * 2 + 1, mid + 1, end, left, right, val);
}
void Game_Start() {
	Dfs(1,-1);
	Init_Tree(1, 1, N);
	ll ans = 0;
	for (int i = 0; i < M; i++) {
		ll a, b;
		cin >> a >> b;
		ans += Query(1, 1, N, index[a].first , index[a].second, b);
		ans %= 1000000007;
	}

	cout << ans;
}

int main(void) {
	ios::sync_with_stdio(false), cin.tie(0), cout.tie(0);
	Input();
	Game_Start();
}