BOJ 25685 - 좋은 노드 집합 찾기 [C++]

 

 

 

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

25685번: 좋은 노드 집합 찾기

문제 요약

정점에 $1$부터 $N$까지 번호가 부여된 루트 있는 트리에서, "좋은 노드 집합"은 이 트리의 노드의 부분집합 $S$이면서 아래 조건을 모두 만족하는 집합으로 정의한다.

1. 어떤 노드 $x$가 $S$에 속해있다면 $x$와 연결된 노드들은 $S$에 속하지 않는다.
2. 어떤 노드 $x$가 $S$에 속하지 않았다면 $x$의 자식 노드들 중 최소 하나의 노드는 $S$에 속한다.

각 노드에 가중치가 있을 때, $S$에 속한 노드의 가중치 합으로 가능한 최댓값을 구하라.

 

문제 해결

이 문제는 전체 문제(트리 전체)를 부분 문제(서브트리)로 쪼개어 생각해도 같은 풀이법을 적용할 수 있는 문제이므로, dp를 적용할 수 있다.

 

$dp1[i]$를 다음과 같이 정의하자.

$i$번 노드를 루트로 하는 서브트리에서, $i$번 노드를 $S$에 넣었을 때의 최대 가중치 합

 

$dp1[i]$는 간단하게 구할 수 있다.

$i$번 노드의 모든 자식들 $c$에 대해서, $dp1[i] = \sum{dp2[c]} + weight[i]$이다.

 

 

$dp2[i]$도 비슷하게 정의하자.

$i$번 노드를 루트로 하는 서브트리에서, $i$번 노드를 $S$에 넣지 않았을 때의 최대 가중치 합

 

$i$번 노드의 모든 자식들 $c$에 대해서, 일반적으론 $dp2[i] = \sum{\max(dp1[c], dp2[c])}$ 이다.

하지만 가중치가 음수일 수 있기 때문에, "적어도 하나의 $c$에 대해선 $S$에 속해야 한다"라는 조건에 어긋날 수 있다.

 

위의 식을 돌렸을 때 모두 dp2의 값만 선택했을 경우 따로 예외처리 해주면 된다.

어떤 한 자식 $c$에 대해서, 반드시 손해를 감수하고 dp1의 값을 선택해야 한다.

이 손해를 최소화하는 방향으로, 즉 $dp1[c] - dp2[c]$가 최대인 한 곳에 대해서만 구해놓은 $dp2[i]$에서 해당 값을 더해주기만 하면 된다.

 

코드

#include <bits/stdc++.h>
using namespace std;
using ll = long long;

vector<vector<int>> V;
vector<ll> arr;
vector<pair<ll, ll>> dp;

void dfs(int n) {
    dp[n].first = arr[n];
    int child = 0, sel = 0, id = 0;
    ll mx = -1e18;
    for (int i : V[n]) {
        dfs(i);
        child++;
        dp[n].first += dp[i].second;
        if (dp[i].first >= dp[i].second) {
            dp[n].second += dp[i].first;
            sel++;
        }
        else {
            ll val = dp[i].first - dp[i].second;
            if (val > mx)	mx = val, id = i;
            dp[n].second += dp[i].second;
        }
    }
    if (child && !sel)	dp[n].second += mx;
}

int main() {
    cin.tie(0)->sync_with_stdio(0);

    int T;
    for (cin >> T; T--;) {
        int n;
        cin >> n;
        V = vector<vector<int>>(n + 1);
        arr = vector<ll>(n + 1);
        dp = vector<pair<ll, ll>>(n + 1);
        int root;
        for (int i = 1; i <= n; i++)	cin >> arr[i];
        for (int i = 1; i <= n; i++) {
            int a;
            cin >> a;
            if (!a)	root = i;
            else	V[a].push_back(i);
        }
        dfs(root);
        cout << max(dp[root].first, dp[root].second) << '\n';
    }

}