[BOJ 25638 // C++] 트리와 경로 개수 쿼리

※ 글쓴이는 취미로 코딩을 익혀보는 사람이라 정확하지 않은 내용을 담고 있을 수 있다 ※

 

이번에 볼 문제는 백준 25638번 문제인 트리와 경로 개수 쿼리이다.
문제는 아래 링크를 확인하자.

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

25638번: 트리와 경로 개수 쿼리

편의상 주어진 트리를 1번 노드를 root로 갖는 rooted tree로 생각하자. 이 때, 각 노드 n에 대하여 해당 노드를 root로 갖는 subtree들을 각각 생각할 수 있다.

 

노드 n을 지나는 경로들은 (1) n을 루트로 갖는 subtree 안에 포함되어있거나 (2) n을 루트로 갖는 subtree에 포함되어있지 않은 두 가지 경우로 나누어 생각할 수 있다. 각 경우의 경로의 수를 각각 구해 문제를 해결해보자.

 

문제 조건에 따라 한 경로가 노드 n을 경유한다는 것은 n이 그 경로의 시작점과 끝점이 아니어야한다는 점을 놓치지 말자.

또한, 답이 32비트 정수로 표현할 수 있는 범위를 넘을 수 있음에 유의하자.

 

아래는 제출한 소스코드이다.

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

int N, Q;
vector<int> G[100001];
bool color[100001]; // 1: red, 0: blue
ll ans[100001];
ll totalred = 0, totalblue = 0;

pair<ll, ll> dfs(ll cur, ll par) { //자신의 부모노드에게 해당 노드를 루트로 하는 subtree에 포함된 각 색의 노드개수 리턴
	ll red = 0, blue = 0;
	
	for (auto& node : G[cur]) {
		if (node == par) continue;
		auto p = dfs(node, cur);
		ans[cur] += red * p.second + blue * p.first;
		red += p.first, blue += p.second;
	}

	if (color[cur]) ans[cur] += red * (totalblue - blue) + blue * (totalred - red - 1);
	else ans[cur] += red * (totalblue - blue - 1) + blue * (totalred - red);

	if (color[cur]) red++;
	else blue++;

	return make_pair(red, blue);
}

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

	cin >> N;
	for (int i = 1; i <= N; i++) {
		cin >> color[i];
		if (color[i]) totalred++;
		else totalblue++;
	}
	for (int i = 1; i < N; i++) {
		int x, y; cin >> x >> y;
		G[x].emplace_back(y);
		G[y].emplace_back(x);
	}

	dfs(1, -1);

	cin >> Q;
	while (Q--) {
		int x; cin >> x;
		cout << ans[x] << '\n';
	}
}