[백준] 등산 마니아 20188번 (c++)

출처 : https://www.acmicpc.net/problem/20188

 

풀이 방법

완전 탐색으로 노드를 기준으로 하면 시간 복잡도가 N^2 이므로 시간 초과된다.

시간 복잡도를 줄이기 위해 DP를 사용할까 했지만, 노드를 기준으로 문제를 완벽하게 분리하는 것이 쉽지 않았다.

이럴 때 필요한 것이 역발상 즉 간선을 기준으로 생각해 보면 정답을 구할 수 있다.

 

위 그림에서 5-6을 잇는 간선을 지나가는 경우는 5를 루트로 하는 서브트리의 노드 중 2개를 선택하거나, 5의 서브트리 노드 중 하나를 선택하고, 1, 4, 3중에 하나를 선택하는 경우라고 알 수 있다.

그러므로, 각 노드를 루트로 가지는 서브트리에서 노드의 수를 계산해 주면 된다.

 

이는 dfs로 계산하고, 나온 값이 ret라고 하면, ret * (ret-1) / 2 + ret * (n - ret)가 루트 노드의 부모로 가는 간선이 선택되는 개수임을 알 수 있다. 

 

주의 점 1일때는 계산하지 않는다. 1은 부모 노드가 없기 때문에

또한 N은 최대 3 * 10 ^ 5 이므로 long long 타입을 사용해야 한다.

 

#include <iostream>
#include <vector>
#include <memory.h>

using namespace std;

int n;
int child_cnt[300001];
vector <vector<int> > graph;
int visited[300001];
long long ans = 0;

int dfs(int idx) { // 자식의 수 세기
	visited[idx] = 1;	
	int ret = 0;

	for (int i = 0; i < graph[idx].size(); i++) {
		if (visited[graph[idx][i]] == 0)
			ret += dfs(graph[idx][i]); // 자식의 수 세기
	}    
    ret += 1;
	if (idx != 1) {
        ans += (static_cast<long>(ret) * (ret - 1) / 2);
        ans += (ret * static_cast<long>(n - ret));
    }
	return ret;
}
int main() {
	cin >> n;
	graph.resize(n + 1);
	
	int a, b;
	for (int i = 0; i < n-1; i++) { // 트리 만들기
		cin >> a >> b;
		graph[a].push_back(b);
		graph[b].push_back(a);
	}
	memset(visited, 0, sizeof(visited));
	dfs(1);	
	cout << ans;

	return 0;
}