CF 1187E Tree Painting

문제

https://codeforces.com/contest/1187/problem/E

트리에서 어떤 노드를 시작으로 선택한 노드들과 인접한 노드들을 골라 검정색으로 칠하는 과정을 계속 반복한다. 이때 노드 하나를 선택할 때 아직 색칠되지 않은 컴포넌트의 크기만큼의 점수를 얻을 때, 점수의 최댓값을 구하는 문제이다.

 

풀이

가장 먼저, 처음에 선택하는 노드를 하나 정해놓고 생각해보자. 다음에 선택하는 노드는 그 노드와 연결되어 있는 점이어야 하고, 다음에 선택할 점들을 어떤 순서로 골라도 상관이 없다.

이를 이용하면, 다음을 알 수 있다.

루트 하나를 고정시켜 놓은 트리에서, 각 노드마다 서브트리의 크기를 구하면 그 서브트리 크기의 합이 바로 답이 된다.

 

이는 트리 DP를 이용하여 구할 수 있다. dp[v]=dp[chd]의 합 + v의 서브트리의 크기

하지만 만약 이 과정을 모든 노드를 루트로 잡았을 때에 대하여 N번 반복한다면 $O(N^2)$ 의 시간이 걸린다.

 

이러한 유형의 문제가 2018년 선발고사 Day 1 1번에 나왔었다.

이와 같이 루트를 정하면 tree dp 를 이용해 풀리는 문제에 대하여 "Rerooting" 이라는 방법을 사용해 볼 수 있다.

 

지금 루트가 u일 때 다음 루트 v로 이동했을 때 dp, sz(서브트리의 사이즈) 배열들을 빠르게 갱신할 수 있으면 된다.

그럼, u에서 u의 자식 v로 이동했다면 어떤 값들이 바뀌게 되는지 확인해 보자.

오직 dp[u], dp[v], sz[u], sz[v] 만 바뀌게 됨을 알 수 있다.

 

이 4개의 dp값들을 바꿔 주는 것은 상수 시간 안에 해결할 수 있고, dfs를 통해 돌면서 dp값들을 갱신시키고, 다시 돌아와서는 원래의 상태로 맞추어 주는 과정을 계속 반복하면 결국 이 문제를 해결할 수 있다.

 

시간 복잡도 : $O(N)$

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

typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

const int MAXN = 2e5;

int N;
ll ans;
vector<int> adj[MAXN+10];

ll dp[MAXN+10], sz[MAXN+10];

void dfs1(int now, int bef)
{
    sz[now]=1;
    for(int nxt : adj[now])
    {
        if(nxt==bef) continue;
        dfs1(nxt, now);
        sz[now]+=sz[nxt];
        dp[now]+=dp[nxt];
    }
    dp[now]+=sz[now];
}

void dfs2(int now, int bef)
{
    ans=max(ans, dp[now]);
    for(int nxt : adj[now])
    {
        if(nxt==bef) continue;

        dp[now]-=dp[nxt]+sz[nxt];
        sz[now]-=sz[nxt];
        sz[nxt]+=sz[now];
        dp[nxt]+=dp[now]+sz[now];

        dfs2(nxt, now);

        dp[nxt]-=dp[now]+sz[now];
        sz[nxt]-=sz[now];
        sz[now]+=sz[nxt];
        dp[now]+=dp[nxt]+sz[nxt];
    }
}

int main()
{
    int i, j;

    scanf("%d", &N);
    for(i=1; i<N; i++)
    {
        int u, v;
        scanf("%d%d", &u, &v);
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    dfs1(1, 1);
    dfs2(1, 1);
    printf("%lld", ans);
}