BOJ 15647 로스팅하는 엠마도 바리스타입니다

백준 15647 링크

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

15647번: 로스팅하는 엠마도 바리스타입니다

풀이

상당히 오래전에 DFS를 N번 도는걸 생각을 했지만 그렇게 해서는 통과가 되지 않는다는게 자명한 사실이므로 그냥 포기한 상태로 그렇게 묻어두다가 최근 앉아서 오래동안 생각할 시간을 얻게 되서 하루종일 수첩을 들고 생각을 해봤음.

 

일단 DFS를 돌리는 횟수를 줄여야 된다는걸 중점으로 생각을 하던 중, 한 노드의 거리(가중치)합을 알게 된다면, 인접한 노드는 간단한 식을 통해서 구할 수 있다는 사실을 발견하였다.

 

트리를 간단히 정리해 다음과 같은 모양이 있다고 하자

위 그림중 A노드의 거리 합을 알고 있다고 하고 그 값을 Sum(A), B노드의 경우는 Sum(B)라고 하고 그 값을 구해보자.

여기서 A와 B노드의 자신을 포함한 자식 노드들의 수를 D(A), D(B)라고 해보자.

그러면 다음과 같은 식을 얻을 수 있다.

 

Sum(B) = Sum(A) - w*( D(B) - D(A) )

 

그런데 여기서 D(A) = N - D(B) 로 대입할 수 있다. 그러면 식을 정리해서

 

Sum(B) = Sum(A) - w*( 2 * D(B) - N )

 

을 얻을 수 있다. 그러면 DFS를 돌리는 횟수를 크게 줄일 수 있다.

 

본인은 DFS를 총 3번 돌렸는데, 한곳을 기준으로 자식노드의 수를 확인할때 한번, 가중치의 합을 구할때 한번, 마지막으로 해당 식을 이용해 각 노드별로 가중치의 합 계산해줄때 한번했다.

 

오버플로우이므로 자료형은 long long int

 

더 줄이고 다듬을 수 있을거같은데 나중에 해봐야겠음 휴가중이라 시간아까워

코드

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

struct edge
{
	int v, w, des;
};

int N, visit[300001];
long long W[300001];
vector <edge> con[300001];

void visit_reset() {
	for (int i = 0; i < 300001; i++)
		visit[i] = 0;
}

int DFS_d(int v) {
	int result = 0;
	visit[v] = 1;
	if (con[v].size() == 1 && visit[con[v][0].v] == 1) {
		//막다른 길
		return 0;
	}
	else {
		for (int i = 0; i < con[v].size(); i++) {
			edge& t = con[v][i];
			if (!visit[t.v]) {
				t.des += 1;
				t.des += DFS_d(t.v);
				result += t.des;
			}
		}
		return result;
	}
}

long long int DFS_w(int v) {
	long long int result = 0;
	visit[v] = 1;
	if (con[v].size() == 1 && visit[con[v][0].v] == 1) {
		//막다른 길
		return 0;
	}
	else {
		for (int i = 0; i < con[v].size(); i++) {
			edge t = con[v][i];
			if (!visit[t.v]) {
				result += t.w*t.des;
				result += DFS_w(t.v);
			}
		}
		return result;
	}
}

int DFS_Wset(int v) {
	visit[v] = 1;
	if (con[v].size() == 1 && visit[con[v][0].v] == 1) {
		//막다른 길
		return 0;
	}
	else {
		for (int i = 0; i < con[v].size(); i++) {
			edge t = con[v][i];
			if (!visit[t.v]) {
				W[t.v] = W[v] - t.w*(2 * t.des - N);
				DFS_Wset(t.v);
			}
		}
		return 0;
	}
}

int main() {
//	freopen("input.txt", "r", stdin);

	cin >> N;
	for (int i = 1; i < N; i++) {
		int u, v, w;
		scanf("%d %d %d", &u, &v, &w);
		con[u].push_back({ v,w,0 });
		con[v].push_back({ u,w,0 });
	}

	DFS_d(1);
	visit_reset();
	W[1] = DFS_w(1);
	visit_reset();
	DFS_Wset(1);

	for (int i = 1; i <= N; i++) {
		printf("%lld\n", W[i]);
	}
}