JUNGOL #4727 두 개의 팀 c++

https://jungol.co.kr/problem/4727

 

에디토리얼을 참고하였음에도 여러번의 시도 끝에 겨우 맞춘 문제이다.

문제

 

음수가중치가 포함된 N개의 노드로 이루어진 트리가 있다. 이 트리에서 두개의 노드를 뽑아 팀장으로 만들 것이다.

팀장이 된 노드는 각각 해당 노드를 루트로 하는 트리에서 가중치 합이 최대가 되도록 노드를 구성해야 한다. 즉, 팀장 노드를 정하는 순간 해당 팀의 트리 모양은 하나로 정해진다. 또한, 두팀의 노드는 서로 겹치면 안된다. 이때, 두개의 팀을 적절히 구성하여 두 팀의 가중치 합의 최댓값을 구하시오.

 

(2 <= N <= 200000)

 

문제 이해

 

문제가 매우 난해하게 느껴지는데, 서브테스크를 하나씩 따라가보며 어떤 흐름으로 만들어졌고, 정확히 무엇을 요구하는 문제인지 차차 알아가보자.

#1	17점	모든 사원의 실력은 양수이다.
#2	12점	N ≤ 5,000 이고, 사원 1을 제외한 각 사원 i의 직속 상사 번호는 i-1이다.
#3	20점	사원 1을 제외한 각 사원 i의 직속 상사 번호는 i-1이다.
#4	16점	N ≤ ​400.
#5	17점	N ≤ ​5,000.
#6	18점	추가 제약 조건 없음.

 

 

우선 모든 노드의 가중치가 양수일 경우를 생각해보자. 각각의 팀장은 트리의 가중치 합을 최대화 해야 하기 때문에, 각각의 팀의 트리는 팀장을 루트로 하는 트리 그 자체일 것이다.

 

문제에서  '두 팀의 노드는 겹치면 안된다'라는 조건을 고려해봤을 때, 모든 노드가 양수일 때에는 두 팀장은 반드시 조상 - 자손관계가 아니어야 한다. 

 

때문에, 모든 노드가 양수일때에는, 두 팀이 최대한 많은 노드를 확보하면서, 두 노드가 조상-자손관계가 아닌 경우는 오직 두개의 트리가 전체 트리의 루트 노드의 두 자매노드일 때 뿐이다. 

이때에는 루트노드의 자식 노드중 트리의 크기가 가장 큰 2개의 자손만을 고르면 된다.

하지만 우리는 음수가중치가 포함된 경우에 대해 생각해보아야 한다.

 

 

문제 접근

모든 노드가 양수일 때와는 달리 이제 두개의 팀은 조상 - 자손 관계에 있을 수 있다.

상위에 있는 팀에서 최적으로 팀을 구성하기 위해서 제외한 노드들로 다시 팀을 최적으로 이룰 수 있기 때문이다.

 

그 다음에는 무얼 해야 할까...

 

우리는 이제 문제를 풀기 위한 "가정"을 해야 한다. 

트리 문제의 기본은 재귀적 연산이라는 것을 잊어서는 안된다.

자식 노드에 대한 "가정"을 미리 해둔 다음, 그에 따라 상위노드의 정보를 채워나가야 하기 때문이다.

 

우리의 가장 첫번째 목표는 "음수가중치를 고려하여 한 팀을 최적으로 구성하는 것"이다. 

 

두개의 팀이 겹치지 않도록 각각 고르는 것은 그 이후의 일이다.

 

1. 최적으로 구성하기 

우선 자식노드들이 모두 최적으로 구성된 팀이라고 해보자.

현재 노드를 팀장으로 하여 팀을 구성한다고 하면, 자식노드들에 대한 검사를 할 것이다.

 

그런데, 생각해보면 우리는 현재 노드가 팀장이라는 가정을 이미 해버렸다. 

자식 노드들의 상태 중 어떤 것을 참조할지도 생각해보자.

 

자식 노드 또한 자식노드가 팀장이라고 가정한 후, 최댓값을 구했을 것이다.

우리는 이 값을 그대로 가져와도 무관하다.

 

애초에 문제 내에서 "팀장"은 "해당 노드로부터 시작해서 합이 최대가 되도록 구성해야 한다" 라는 조건이기 때문에, 

그냥 현재 노드에서 시작하는 최적 트리를 저장하는 값으로 치환하여 생각해도 무관하기 때문이다.

 

i번 노드가 팀장일 때 팀원의 가중치의 최대합을 team[i]라고 하자.

 

이러한 흐름에 따라 현재 노드의 자식 노드들 c의 team[c] 중 양수인 것 만을 현재 노드의 가중치와 더하여

team[cur]를 구해낼 수 있다.

 

2. 두개의 팀을 합치는 첫번째 방법

 

그런데, 앞서 우리가 구한 team[i]는 i번째 노드를 팀장으로 해야 한다는 제약이 존재한다.

때문에, 팀의 최적합 중 "최신 최댓값" 을 구하기 위해서는 이전까지의 team[i]값  중 최댓값을 저장하는 prefix_max 배열을 하나 더 만들어야 한다. 

나는 이 배열을 sub[i]로 잡았다.

이렇게 되면, 어떤 노드에 대해서 자식 노드 c에 대해서 sub[c] 중에서 값이 가장 큰 2개를 골라서 더해주면, 

최소공통조상이 현재 노드인 두개의 팀의 총합 중 최댓값을 구할 수 있는 것이다.

 

 

 

 

3. 두개의 팀을 합치는 두번째 방법

 

이제 팀장이 조상-자손 관계에 있을 때를 생각해보아야 한다.

이제 우리는 team과 sub의 정의를 자유롭게 사용할 수 있다.

 

앞서 우리가 team[cur]를 구성할 때 양수가 아닌 team[c]를 제외했던 것을 기억할 것이다. 

제외되었다는 것은 team[cur]의 팀과 별개라는 것이다.

i 팀에서 제외된 노드 중 최대합 exc[i] 라는 배열에 저장하자.

 

제외된 c의 서브트리 중에서 최대인 팀 구성, 즉 sub[c]와 비교하여 exc[cur]를 업데이트해줄 수 있다.

 

또한, 노드를 포함하는 경우에도 선택한 team[c] 중에서 제외된 노드들 중 최대 트리 합 exc[c]와도 비교하여 이전 최댓값을 

유지할 수도 있다.

 

max(다른 팀의 노드 중 하위 노드의 제외된 최대 트리합, 선택되지 않은 노드의 최대 트리 합)

 

 

이렇게 되면 매번 exc가 제대로 업데이트 된다.

(주의해야 할점은 선택하는 기준은 반드시 '양수인지의 여부'라는 것이다. 최대합이 0인 부분은 최적 선택 조건 상 선택해도 되고 선택하지 않아도 그만이지만, team1이 이유 없이 점유해버릴 경우에는 team2가 sub[c]로 접근할 수 있는 기회가 사라져 exc[cur]의 값이 작아지기 때문이다.)

 

결국 team[i] + exc[i]의 합을 통해 두 팀의 합을 구하는 두번째 방법도 가능해진 것이다.

 

방법1과 방법2를 dfs로 순서대로 적용시켜 그 중 최댓값을 찾아주면 정답이다.

 

 

 

정답코드

 

#include <bits/stdc++.h>
#define int long long
#define MIN -10e17
using namespace std;
vector<int>w, team, sub, exc;
vector<vector<int>>g;
int ans = MIN;
void dfs(int p, int cur) {
	for(int c: g[cur]) {
		if(c==p)continue;
		dfs(cur, c);
	}
	team[cur] = max(team[cur], w[cur]);
	for(int c: g[cur]) {
		if(c==p)continue;
		sub[cur] = max(sub[cur], sub[c]);
		if(team[c] <= 0) {
			exc[cur] = max(exc[cur], sub[c]);
			continue;
		}
		team[cur] += team[c];
		exc[cur] = max(exc[cur], exc[c]);
	}
	sub[cur] = max(sub[cur], team[cur]);
	ans = max(ans, team[cur]+exc[cur]);
}

void dfs2(int p, int cur) {
	for(int c: g[cur]) {
		if(c==p)continue;
		dfs2(cur, c);
	}
	int nc = 0;
	priority_queue<int> pq;
	for(int c : g[cur]) {
		if(c==p)continue;
		pq.push(sub[c]);
		nc++;
	}
	if(nc < 2)return;
	int k = 0;
	k += pq.top();
	pq.pop();
	k += pq.top();
	ans = max(ans, k);
}

int32_t main() {
	int n;
	cin >> n;
	w.resize(n+1);
	g.resize(n+1);
	team.resize(n+1, MIN);
	sub.resize(n+1, MIN);
	exc.resize(n+1, MIN);
	int a, b;
	cin >> a >> b;
	w[1] = a;
	for(int i = 2; i <= n; i++) {
		int x;
		cin >> w[i] >> x;
		g[i].push_back(x);
		g[x].push_back(i);
	}
	dfs(-1, 1);
	dfs2(-1, 1);
	cout << ans;
}

 

 

 

 

+

어려워도 너무 어렵다. 이렇게 복잡하게 얽히고 섥혀 있는 관계가 사람의 머릿속에서 나올 수 있다는 것이 신기하다.

트리 dp의 문제 접근 방식은 아직 많이 연습해야 하는 부분인 것 같다.