[KOI] 백준 21759 - 두 개의 팀

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

 

이 포스트에서는 문제 지문을 완벽히 이해했다고 가정하고 풀이를 설명한다.

 

꽤 까다로운 트리 DP이다. 기본적인 트리 DP만 해봤다면 상당히 애를 먹을 수도 있다. 우선 다음 3개의 테이블이 필요하다.

 

dp1(i) := i가 팀장이 되어 팀을 구성했을 때 최대 점수

 

dp2(i) := i가 팀장일 때 i의 자식 노드들에 대해 dp1(i)에 속하지 않은 정점을 대상으로 팀을 구성했을 때 최대 점수

 

dp3(i) := i가 루트인 서브 트리가 가질 수 있는 최대 점수

 

각 테이블을 계산해보자.

 

1. dp1을 계산하는 것은 쉽다. dfs 재귀로 합 dp를 수행하면 된다. 이때 자식 노드의 dp1 값이 음수인 경우는 더하지 않는 게 최적이다.

 

2. 자식 노드 i에 대해 $dp1(i) \leq 0$인 경우 dp1을 계산하지 않고 현재 노드를 x라 할 때 $dp2(x) = max(dp2(x), dp3(i))$를 계산하다. dp1(i)를 더하지 않기로 했으므로 dp2(x)의 계산에 dp3(i)로 파고 들어가는 과정이 있는 것은 당연하다.

 

3. 자식 노드 i에 대해 $dp1(i) > 0$인 경우 현재 노드가 x라 할 때 $dp1(x) = dp1(x) + dp1(i)$로 값을 경신하고 $dp2(x) = max(dp2(x), dp2(i))$를 계산한다. 이는 dp1에 i가 편입되었으므로 i를 제외한 팀의 최댓값을 찾기 위해 내려 보내는 것으로 이해할 수 있다.

 

1~3의 과정을 거치면서 자식 노드 i들의 dp3(i) 값 중 최댓값 2개를 보존한다. 이 포스트의 구현에서는 max heap을 사용할 것이다. 이는 현재 노드 x가 팀장일 때 최적해가 아닌 경우 자식 서브 트리들에게서 최적해를 도출하기 위함이다.

 

자식 노드 i들에게서 위 테이블 계산을 끝내면 dp1(x)와 dp2(x)의 합의 최댓값으로 최적해의 후보를 계산할 수 있다.

 

그리고 트리의 위층 노드들을 계산하기 위해 dp3 테이블도 계산해야 한다. 우선 dp3(x) = dp1(x)로 값을 한 번 초기화 해준 다음에 max heap에 넣은 노드 중 최댓값을 기존 dp3(x)와 비교 및 대입하면 된다.

 

그리고 아까 언급했던 자식 서브트리들에게서 최적해가 나오는 경우를 고려하기 위해 정답을 max heap에 있던 가장 큰 2개의 값의 합을 반영해주면 최종적으로 답을 구할 수 있다.

 

전체 코드

#pragma GCC target("sse,sse2,sse3,ssse3,sse4,avx,avx2,fma")
#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")
 
#include "bits/stdc++.h"
#include "ext/rope"
 
using namespace std;
using namespace __gnu_cxx;
 
using ll = long long;
using pii = pair<int, int>;
 
int n;
int skill[200001];
ll dp1[200001], dp2[200001], dp3[200001];
vector<int> tree[200001];
ll ans = -1e18;
 
void dfs(int cur)
{
  dp1[cur] = skill[cur];
  dp2[cur] = -1e18;
 
  priority_queue<ll> pq;
 
  for (int i : tree[cur])
  {
    dfs(i);
    if (dp1[i] > 0)
    {
      dp1[cur] += dp1[i];
      dp2[cur] = max(dp2[cur], dp2[i]);
    }
    else
    {
      dp2[cur] = max(dp2[cur], dp3[i]);
    }
    pq.push(dp3[i]);
  }
 
  dp3[cur] = dp1[cur];
  ans = max(ans, dp1[cur] + dp2[cur]);
 
  while (pq.size())
  {
    auto now = pq.top();
    pq.pop();
 
    dp3[cur] = max(dp3[cur], now);
    if (pq.size()) ans = max(ans, now + pq.top());
    return;
  }
}
 
void solve()
{
  cin >> n;
  for (int i = 1; i <= n; i++)
  {
    cin >> skill[i];
    int par;
    cin >> par;
    if (par != -1) tree[par].push_back(i);
  }
 
  dfs(1);
  cout << ans;
}
 
int main()
{
#ifdef LOCAL
  freopen("input.txt", "r", stdin);
//  freopen("output.txt", "w", stdout);
#endif
  cin.tie(nullptr);
  ios::sync_with_stdio(false);
 
  solve();
}

제출 기록