[백준/BOJ] 백준 15669번 : 나무 위의 입자

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

15669번: 나무 위의 입자

1번 정점이 루트인 트리를 만들고, cache[100001][2]에 해당 정점이 루트인 서브트리에서 깊이가 홀수인 노드의 개수와 깊이가 짝수인 노드의 개수를 저장해 놓고, 도착점이 빨간색일 때 이동하는 간선의 총개수가 짝수개이므로 (짝수 깊이 노드, 짝수 깊이 노드) 또는 (홀수 깊이 노드, 홀수 깊이 노드)의 조합이어야 하고, 도착점이 검은색일 때는 이동하는 간선의 총개수가 홀수개이므로 (짝수 깊이 노드, 홀수 깊이 노드) 또는 (홀수 깊이 노드, 짝수 깊이 노드)의 조합이어야 한다는 것을 이용해 문제를 해결했다.

 

코드

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

int n, m;
vector<int> adj[100001];

vector<int> maked(100001, 0);
vector<int> children[100001]; //트리에서 자식노드
vector<int> depth(100001, 0);

int cache[100001][2]; //현재 위치가 루트인 서브트리에서 깊이가 홀수인 노드의 개수와 깊이가 짝수인 노드의 개수 

void Pre()
{
	for (int i = 0; i < 100001; i++)
	{
		for (int j = 0; j < 2; j++)
		{
			cache[i][j] = 0;
		}
	}
}

void MakeTree(int here, int this_depth)
{
	maked[here] = 1;
	depth[here] = this_depth;

	for (int i = 0; i < adj[here].size(); i++)
	{
		int there = adj[here][i];

		if (maked[there] == 0)
		{
			MakeTree(there, this_depth + 1);
			children[here].push_back(there);
		}
	}
}

//cache 채우기
pair<int, int> Solve(int here)
{
	int& ret1 = cache[here][0];
	int& ret2 = cache[here][1];

	//here의 깊이가 홀수일때
	if (depth[here] % 2 == 1)
	{
		ret1++;
	}

	//here의 깊이가 짝수일때
	else
	{
		ret2++;
	}

	for (int i = 0; i < children[here].size(); i++)
	{
		int there = children[here][i];
		pair<int, int> there_ret = Solve(there);

		ret1 += there_ret.first;
		ret2 += there_ret.second;
	}

	return make_pair(ret1, ret2);
}

int main()
{
	cin.tie(NULL);
	ios_base::sync_with_stdio(false);

	Pre();

	cin >> n >> m;

	for (int i = 0; i < n - 1; i++)
	{
		int u, v;
		cin >> u >> v;

		adj[u].push_back(v);
		adj[v].push_back(u);
	}

	//1번 정점이 루트인 트리를 만든다
	MakeTree(1, 0);

	Solve(1);

	for (int i = 0; i < m; i++)
	{
		int u, v, c;
		long long result = 0;

		cin >> u >> v >> c;

		//v쪽이 u의 부모노드일때 u와 v를 바꾼다(방향을 바꾼다)
		//그러면 트리에서 항상 u가 위쪽에 있게 된다
		if (depth[v] < depth[u])
		{
			swap(u, v);
		}

		//도착점이 빨강일때
		//이동하는 간선의 총 개수가 짝수개여야 한다
		//(짝수깊이노드,짝수깊이노드) 또는 (홀수깊이노드,홀수깊이노드)이어야 된다
		if (c == 0)
		{
			//홀수깊이노드, 홀수깊이노드의 연결일때
			result += ((long long)cache[v][0] * (long long)(cache[1][0] - cache[v][0]));

			//짝수깊이노드, 짝수깊이노드의 연결일때
			result += ((long long)cache[v][1] * (long long)(cache[1][1] - cache[v][1]));
		}

		//도착점이 검정일때
		//이동하는 간선의 총 개수가 홀수개여야 한다
		//(짝수깊이노드,홀수깊이노드) 또는 (홀수깊이노드,짝수깊이노드)이어야 된다
		else
		{
			//짝수깊이노드, 홀수깊이노드의 연결일때
			result += ((long long)cache[v][1] * (long long)(cache[1][0] - cache[v][0]));

			//홀수깊이노드, 짝수깊이노드의 연결일때
			result += ((long long)cache[v][0] * (long long)(cache[1][1] - cache[v][1]));
		}

		cout << result << "\n";
	}

	return 0;
}