[BOJ 14287 | C++] 회사 문화 3 :: ETT(오일러 경로 테크닉)

• 문제 링크 : https://www.acmicpc.net/problem/14287
• 문제 태그 : 자료 구조, 트리, 세그먼트 트리, 오일러 경로 테크닉

 

코드

dfs를 이용하여 in 배열과 out 배열을 관리하고 세그먼트 트리를 활용하여 서브트리의 합을 구했습니다.

 

dfs를 비재귀로 짜는 게 습관이라서 재귀 구현은 쉽지 않았지만 N과 M시리즈를 뽀개며 공부했더니 잘 구현할 수 있었습니다. 세그먼트 트리를 오랜만에 구현했더니 처음에 몇 군데 실수가 있었습니다...😭

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

int main() {
	ios::sync_with_stdio(false);
	cin.tie(NULL);
	cout.tie(NULL);
	
	int n, m;	cin >> n >> m;

	vector<int>in(n + 1);
	vector<int>out(n + 1);
	vector<vector<int>>g(n + 1);

	for (int i = 1;i<=n;i++) {
		int x;	cin >> x;
		if (x == -1) continue;
		g[x].push_back(i);
	}

	int cnt = 0;
	auto recursive = [&](const auto& self, int n) -> void {
		in[n] = ++cnt;
		for (auto x : g[n]) {
			self(self, x);
		}
		out[n] = cnt;
	};
	recursive(recursive, 1);
	

	int sz = 1;
	while(sz < n) sz *= 2;
	vector<int>tree(sz * 2, 0);

	auto update = [&](int i, int w) {
		i = i - 1 + sz;
		tree[i] += w;
		while(i > 1) {
			i /= 2;
			tree[i] = tree[i * 2] + tree[i * 2 + 1];
		}
	};

	auto query = [&](int l, int r) {
		l = l - 1 + sz;
		r = r - 1 + sz;
		
		int res = 0;
		while(l <= r) {
			if (l % 2 == 1)	res += tree[l], l++;
			if (r % 2 == 0) res += tree[r], r--;

			l /= 2;
			r /= 2;
		}

		return res;
	};

	for (int i = 0;i<m;i++) {
		int q, x, w;
		cin >> q >> x;
		if (q == 1) {
			cin >> w;
			update(in[x], w);
		}
		else {
			cout << query(in[x], out[x]) << '\n';
		}
	}
}