[P3] 백준 18227번 성대나라의 물탱크 C++ 세그먼트 트리, 오일러 경로 테크닉

리뷰

 

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

오일러 경로를 구할 때 노드의 깊이를 구해준 후 구간합 쿼리에 해당 깊이를 곱해준 뒤 출력하는 문제

회사 문화 문제와 비슷하나 노드의 깊이를 구해주어야 하는 조건이 추가된 문제

 

 

전역 변수

• N : 배열의 최대 크기를 저장할 상수 변수
• n : 도시의 수를 저장할 변수
• c : 수도의 번호를 저장할 변수
• q : 쿼리의 개수를 저장할 변수
• tree : 세그먼트 트리 정보를 저장할 배열
• it : 오일러 경로의 진입 시점을 저장할 배열
• ot : 오일러 경로의 탈출 시점을 저장할 배열
• dep : 노드의 깊이를 저장할 배열
• t : 오일러 경로의 시간을 저장할 변수
• edges : 인접 리스트를 저장할 벡터 배열

 

함수

1. dfs

void dfs(int level, int node, int par) {
	it[node] = ++t;
	dep[node] = level;
	for (int child : edges[node]) {
		if (child != par) dfs(level + 1, child, node);
	}
	ot[node] = t;
}

 

오일러 경로와 노드의 깊이를 초기화 하기위한 함수

 

2. update

void update(int node, int s, int e, int idx) {
	if (s == e) tree[node]++;
	else {
		int mid = (s + e) / 2;
		if (idx <= mid) update(node * 2, s, mid, idx);
		else update(node * 2 + 1, mid + 1, e, idx);
		tree[node] = tree[node * 2] + tree[node * 2 + 1];
	}
}

 

세그먼트 트리 정보를 업데이트 하기 위한 함수

 

3. query

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

 

세그먼트 트리의 구간합을 구하기 위한 함수

 

문제풀이

1. n, c값을 입력 받고, n - 1개의 간선 정보를 입력 받아 인접 리스트에 양방향으로 추가해 준다.
2. dfs함수를 통해 오일러 경로와 각 노드의 깊이 정보를 it, ot, dep 배열에 저장해 준다.
3. q값을 입력 받고, q개의 쿼리의 인자를 각각 변수 a, b에 입력 받아준다.
4. a가 1일 경우 update함수의 세그먼트 트리의 it[b]에 해당하는 인덱스 값을 1만큼 증가시킨다.
5. a가 2일 경우 query함수를 통해 세그먼트 트리의 it[b], ot[b] 구간의 누적합과 dep[b]를 곱한 값을 출력한다.

 

트러블 슈팅

1. 빠른 입출력을 사용하지 않았더니 시간 초과가 발생하였다.
2. 빠른 입출력을 사용해 주니 AC를 받았다.

 

참고 사항

• 루트로 부터의 깊이에 따라 저장되는 물양이 다른 것을 트리에서의 depth로 해결하였다.

 

정답 코드

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

const int N = 2e5 + 1;
int n, c, q;
ll tree[N * 4];
int it[N];
int ot[N];
int dep[N];
int t;
vector<int> edges[N];

void dfs(int level, int node, int par) {
	it[node] = ++t;
	dep[node] = level;
	for (int child : edges[node]) {
		if (child != par) dfs(level + 1, child, node);
	}
	ot[node] = t;
}

void update(int node, int s, int e, int idx) {
	if (s == e) tree[node]++;
	else {
		int mid = (s + e) / 2;
		if (idx <= mid) update(node * 2, s, mid, idx);
		else update(node * 2 + 1, mid + 1, e, idx);
		tree[node] = tree[node * 2] + tree[node * 2 + 1];
	}
}

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

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

	cin >> n >> c;
	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, c, -1);
	//for (int i = 1; i <= n; ++i) cout << i << " : " << it[i] << " " << ot[i] << "\n";
	cin >> q;
	while (q--) {
		int a, b; cin >> a >> b;
		if (a == 1) update(1, 1, n, it[b]);
		else cout << query(1, 1, n, it[b], ot[b]) * dep[b] << "\n";
	}
}