Maximum White Subtree

문제

링크: Maximum White Subtree

난이도: 1800

풀이

정점 $v$ 를 포함하는 서브트리의 $cnt_w - cnt_b$ 의 최댓값을 $ans_v$ 라 합시다.

전체 문제를 풀기 전에 앞서 트리의 루트 $r$ 에 대해 $ans_r$ 을 구해 봅시다.

이는 트리 DP로 재귀적으로 구할 수 있습니다.

정점 $v$ 를 루트로 하는 서브트리만 고려했을 때 $v$ 를 포함하는 서브트리의 $cnt_w - cnt_b$ 의 최댓값을 $dp_v$ 라 합시다.

$v$ 가 리프 노드인 경우 $v$ 의 색에 따라서 $dp_v$ 가 정해집니다.

리프 노드가 아닌 경우 $v$ 의 자식 $c$ 에 대해 $dp_c > 0$ 인 경우를 모두 서브트리에 포함시키면 $cnt_w - cnt_b$ 가 최대가 되므로 $dp_v = \sum \max(0, dp_c)$ 입니다.

해당 점화식으로 $dp_r$ 을 $O(n)$ 에 구할 수 있습니다.

$r$ 을 루트로 하는 서브트리는 전체 트리와 같으므로 $ans_r = dp_r$ 입니다.

이제 전방향 DP를 이용해서 $r$ 의 한 자식 $c$ 에 대한 $ans$ 를 구할 수 있습니다.

먼저 $c$ 가 $r$ 의 정답 서브트리에 포함되는 경우를 생각해 봅시다. 그렇다면 $dp_c > 0$ 이어야 합니다.

$r$ 을 포함하는 경우는 자명하게 $ans_r$ 과 같습니다.

$r$ 을 포함하지 않는 경우는 $c$ 의 서브트리로만 구성하는 경우와 같고, 따라서 $dp_c$ 가 됩니다.

따라서 $ans_c = \max(dp_c, ans_r)$ 입니다.

$c$ 가 $r$ 의 정답 서브트리에 포함되지 않는 경우를 생각해 봅시다. 그렇다면 $dp_c \leq 0$ 입니다.

$r$ 을 포함하는 경우는 $dp_c + ans_r$ 입니다.

$r$ 을 포함하지 않는 경우는 위에서 본 것과 같이 $dp_c$ 가 됩니다.

따라서 $ans_c = \max(dp_c, dp_c + ans_r)$ 입니다.

구한 $ans_c$ 를 이용해서 $c$ 의 자식들의 $ans$ 도 계산해주면 모든 정점의 $ans$ 를 총 $O(n)$ 에 계산할 수 있습니다.

코드

#include <bits/stdc++.h>
using namespace std;
 
int color[202020];
vector<int> adj[202020];
 
int subdiff[202020];
int ans[202020];
 
void calcDiff(int cur, int parent) {
    subdiff[cur] = color[cur];
 
    for (int child : adj[cur]) {
        if (child == parent) continue;
 
        calcDiff(child, cur);
        subdiff[cur] += max(0, subdiff[child]);
    }
}
void calcAns(int cur, int parent) {
    for (int child : adj[cur]) {
        if (child == parent) continue;
 
        if (subdiff[child] <= 0)
            ans[child] = max(subdiff[child], subdiff[child] + ans[cur]);
        else ans[child] = max(subdiff[child], ans[cur]);
 
        calcAns(child, cur);
    }
}
 
int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
 
    int n;
    cin >> n;
 
    for (int i = 1; i <= n; i++) {
        int a;
        cin >> a;
 
        color[i] = 2 * a - 1;  // 1 if white and -1 if black
    }
 
    for (int i = 0; i < n - 1; i++) {
        int u, v;
        cin >> u >> v;
 
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
 
    calcDiff(1, 0);
    ans[1] = subdiff[1];
    calcAns(1, 0);
 
    for (int i = 1; i <= n; i++) cout << ans[i] << ' ';
 
    return 0;
}

문제

링크: Maximum White Subtree

난이도: 1800

풀이

정점 $v$ 를 포함하는 서브트리의 $cnt_w - cnt_b$ 의 최댓값을 $ans_v$ 라 합시다.

전체 문제를 풀기 전에 앞서 트리의 루트 $r$ 에 대해 $ans_r$ 을 구해 봅시다.

이는 트리 DP로 재귀적으로 구할 수 있습니다.

정점 $v$ 를 루트로 하는 서브트리만 고려했을 때 $v$ 를 포함하는 서브트리의 $cnt_w - cnt_b$ 의 최댓값을 $dp_v$ 라 합시다.

$v$ 가 리프 노드인 경우 $v$ 의 색에 따라서 $dp_v$ 가 정해집니다.

해당 점화식으로 $dp_r$ 을 $O(n)$ 에 구할 수 있습니다.

$r$ 을 루트로 하는 서브트리는 전체 트리와 같으므로 $ans_r = dp_r$ 입니다.

이제 전방향 DP를 이용해서 $r$ 의 한 자식 $c$ 에 대한 $ans$ 를 구할 수 있습니다.

먼저 $c$ 가 $r$ 의 정답 서브트리에 포함되는 경우를 생각해 봅시다. 그렇다면 $dp_c > 0$ 이어야 합니다.

$r$ 을 포함하는 경우는 자명하게 $ans_r$ 과 같습니다.

$r$ 을 포함하지 않는 경우는 $c$ 의 서브트리로만 구성하는 경우와 같고, 따라서 $dp_c$ 가 됩니다.

따라서 $ans_c = \max(dp_c, ans_r)$ 입니다.

$c$ 가 $r$ 의 정답 서브트리에 포함되지 않는 경우를 생각해 봅시다. 그렇다면 $dp_c \leq 0$ 입니다.

$r$ 을 포함하는 경우는 $dp_c + ans_r$ 입니다.

$r$ 을 포함하지 않는 경우는 위에서 본 것과 같이 $dp_c$ 가 됩니다.

따라서 $ans_c = \max(dp_c, dp_c + ans_r)$ 입니다.

구한 $ans_c$ 를 이용해서 $c$ 의 자식들의 $ans$ 도 계산해주면 모든 정점의 $ans$ 를 총 $O(n)$ 에 계산할 수 있습니다.

코드

#include <bits/stdc++.h>
using namespace std;
 
int color[202020];
vector<int> adj[202020];
 
int subdiff[202020];
int ans[202020];
 
void calcDiff(int cur, int parent) {
    subdiff[cur] = color[cur];
 
    for (int child : adj[cur]) {
        if (child == parent) continue;
 
        calcDiff(child, cur);
        subdiff[cur] += max(0, subdiff[child]);
    }
}
void calcAns(int cur, int parent) {
    for (int child : adj[cur]) {
        if (child == parent) continue;
 
        if (subdiff[child] <= 0)
            ans[child] = max(subdiff[child], subdiff[child] + ans[cur]);
        else ans[child] = max(subdiff[child], ans[cur]);
 
        calcAns(child, cur);
    }
}
 
int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
 
    int n;
    cin >> n;
 
    for (int i = 1; i <= n; i++) {
        int a;
        cin >> a;
 
        color[i] = 2 * a - 1;  // 1 if white and -1 if black
    }
 
    for (int i = 0; i < n - 1; i++) {
        int u, v;
        cin >> u >> v;
 
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
 
    calcDiff(1, 0);
    ans[1] = subdiff[1];
    calcAns(1, 0);
 
    for (int i = 1; i <= n; i++) cout << ans[i] << ' ';
 
    return 0;
}

Maximum White Subtree

문제

풀이

코드

댓글

Maximum White Subtree

문제

풀이

코드

댓글