JOI19 Unique Cities

문제

https://oj.uz/problem/view/JOI19_ho_t5

문제 보기 - Unique Cities (JOI19_ho_t5) :: oj.uz

트리가 주어지고, 각 정점에 색이 칠해져 있을 때, 정점 $x$에 대한 unique한 정점 $y$를 다음과 같이 정의한다.

$x$와 $y$의 거리가 $d$일 때 $x$에서 거리가 $d$인 $y$가 아닌 다른 정점 $z$가 존재하지 않는다.

각 정점에 대해, unique한 정점들의 서로 다른 색들의 개수를 구해야 한다.

$M<=N<=2*10^5$

 

풀이

Observation 1 : 어떤 정점 $x$에 대한 unique한 정점들은 모두 x에서 가장 멀리 떨어진 정점 $u$까지의 경로 상에 존재한다.

어떤 정점 $x$에 대한 unique한 정점들은 모두 지름을 구성하는 정점 $u$, $v$중 더 멀리 떨어진 정점까지의 경로 상에 존재한다.

최장경로 상에 존재하지 않는 어떤 정점 $v$가 unique 하다고 가정하자. $u$가 가장 멀리 떨어진 정점이기 때문에, $x$, $u$의 경로 상에 dist($u$, $v$)만큼 떨어진 정점이 무조건 존재하니, $v$는 unique할 수가 없어, 모순이다.

가장 멀리 떨어진 정점이 나왔으니, 트리의 지름을 구하는 방법을 떠올려 보면, 다음과 같은 lemma가 성립한다.

트리의 지름을 이루는 두 정점 $u$, $v$중 하나는 각 정점에서 가장 멀리 떨어진 정점이다.

따라서 최장경로를 이루는 정점 $u$는 사실 지름을 구성하는 정점들 중 하나이다.

 

이제 각 정점마다 가장 먼 점을 새로 본다기보다는 고정된 트리의 지름을 구성하는 정점 $u$, $v$에 대하여 탐색을 하며 나머지 정점들에 대한 답을 구하려고 노력해보자.

 

 

전체적인 꼴은 위 그림과 같이 지름이 있고, $u$에서 임의의 정점 $x$까지의 경로를 생각한다고 하면 경로 밖에 있는 삼각형 모양의 서브트리의 정점들은 모두 그 서브트리의 높이만큼 뒤쪽의 정점들이 unique하지 못하도록 만들어주는 역할을 한다. 여기서 서브트리의 구조와 관계 없이 높이만 삭제 정점에 영향을 미친다는 사실에 주의하자.

 

문제의 Subtask 3 (각 정점마다 색이 모두 달라, 단순히 unique한 정점의 수를 세는 문제)부터 해결하자.

$u$를 고정하고, 루트를 $u$로 하는 트리에서 문제를 해결해 보자.

 

기본적으로, 다음과 같은 생각을 할 수 있다.

unique 한 정점들의 후보 집합을 스택으로 관리하고, 하나의 서브트리로 타고 내려갈 때마다 남은 서브트리의 높이만큼 스택에서 제거해 주는 과정을 반복하면 문제를 해결할 수 있다.

하지만 이 방법을 사용하면 다시 올라왔을 때 삭제한 만큼의 정점들을 다시 복구해 주어야 하고, 이런 과정을 거치면 시간 복잡도가 최대 $O(N^2)$까지 올라갈 수 있다.

 

이를 해결하는 아주 간단한 방법이 있다. 바로 가장 큰(높은) 서브트리부터 방문하는 것이다.

Observation 2 : 가장 큰(높은) 서브트리부터 방문한다면, 한번 삭제된 정점이 다시 살아나는 경우는 오직 $x$밖에 없다.

처음에 가장 큰 서브트리로 내려간다는 것은 남은 서브트리의 높이만큼의 점들이 스택에서 제거된다. 그 후 $x$가 추가되고, 가장 큰 서브트리에 의해 정점들이 삭제된다. 다시 $x$로 올라왔을 때, 원래 알고리즘에서는 처음에 삭제한 점들을 복구한 후, 가장 큰 서브트리의 높이만큼의 정점들을 다시 삭제하는 작업을 한 후 $x$를 추가해야 한다.

하지만 가장 큰 서브트리로 들어간 이후 삭제되는 정점들의 최대 높이는 가장 큰 서브트리에 의해 삭제되는 정점보다는 무조건 낮다. 또한, 처음에 가장 큰 서브트리로 들어갔기 때문에 두번째로 큰 서브트리의 높이에 의해 삭제되는 정점들보다 가장 큰 서브트리에 의해 삭제되는 정점의 수가 더 크다. 이제 고려해야 할 점은 오직 $x$밖에 없다. ($x$는 서브트리에 의해 삭제되었어도, 다시 내려갈 때는 추가해 주어야 하기 때문이다.) 하지만 $x$는 최대 자식의 개수만큼 추가, 삭제되며, 트리에서 자식의 개수를 다 더해도 $O(N)$이니, 이 알고리즘을 통해 해결할 수 있다.

위 알고리즘을 트리의 지름 $u$, $v$에 대하여 두번 돌려주면 된다.

 

이제 남은 것은 서로 다른 색의 수를 세는 것 뿐이다.

이를 해결하는 방법은 같은 색이면 루트에 가장 가까운 정점 하나만을 저장하는 것이다. 만약 이 정점이 스택에서 제거되지만 않는다면 이 색을 세어야 하고, 거꾸로 이 정점이 제거된다는 뜻은 스택 상에서 이 정점보다 뒤쪽에 위치한 모든 같은 색의 정점들 또한 제거되었다는 뜻이니, 이 정점의 존재 유무가 결국 스택 전체에서 그 색의 존재 유무를 대변한다고 생각해도 된다.

 

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

 

#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, M, A[MAXN+10];
vector<int> adj[MAXN+10];
 
int P, Q;
int DP[MAXN+10], DQ[MAXN+10];
 
void dfs(int now, int bef, int *D, int dep)
{
    D[now]=dep;
    for(int nxt : adj[now])
    {
        if(nxt==bef) continue;
        dfs(nxt, now, D, dep+1);
    }
}
 
int H[MAXN+10];
vector<int> chd[MAXN+10];
 
void dfs2(int now, int bef)
{
    H[now]=0;
    chd[now].clear();
    for(int nxt : adj[now])
    {
        if(nxt==bef) continue;
        chd[now].push_back(nxt);
        dfs2(nxt, now);
        H[now]=max(H[now], H[nxt]+1);
    }
    sort(chd[now].begin(), chd[now].end(), [&](const int &p, const int &q) { return H[p]>H[q]; });
}
 
int cnt[MAXN+10], chk[MAXN+10], ans[MAXN+10];
vector<pii> V;
 
void dfs3(int now, int bef, int dep)
{
    int i, j;
    if(chd[now].size()==0)
    {
        if(chk[now]) ans[now]=V.size();
    }
    else if(chd[now].size()==1)
    {
        if((V.empty() || V.back().first!=dep) && !cnt[A[now]]) V.push_back({dep, A[now]}), cnt[A[now]]++;
        dfs3(chd[now][0], now, dep+1);
        while(!V.empty() && V.back().first>=dep-H[now]) cnt[V.back().second]--, V.pop_back();
        if(chk[now]) ans[now]=V.size();
    }
    else
    {
        while(!V.empty() && V.back().first>=dep-H[chd[now][1]]-1) cnt[V.back().second]--, V.pop_back();
        if((V.empty() || V.back().first!=dep) && !cnt[A[now]]) V.push_back({dep, A[now]}), cnt[A[now]]++;
        dfs3(chd[now][0], now, dep+1);
        while(!V.empty() && V.back().first>=dep-H[chd[now][0]]-1) cnt[V.back().second]--, V.pop_back();
        for(i=1; i<chd[now].size(); i++)
        {
            int nxt=chd[now][i];
            if((V.empty() || V.back().first!=dep) && !cnt[A[now]]) V.push_back({dep, A[now]}), cnt[A[now]]++;
            dfs3(nxt, now, dep+1);
        }
        while(!V.empty() && V.back().first>=dep-H[now]) cnt[V.back().second]--, V.pop_back();
        if(chk[now]) ans[now]=V.size();
    }
}  
 
int main()
{
    int i, j;
 
    scanf("%d%d", &N, &M);
    for(i=1; i<N; i++)
    {
        int u, v;
        scanf("%d%d", &u, &v);
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    for(i=1; i<=N; i++) scanf("%d", &A[i]);
 
    dfs(1, 1, DP, 0); P=1;
    for(i=1; i<=N; i++) if(DP[P]<DP[i]) P=i;
    dfs(P, P, DP, 0); Q=1;
    for(i=1; i<=N; i++) if(DP[Q]<DP[i]) Q=i;
    dfs(Q, Q, DQ, 0);
 
    //printf("%d %d\n", P, Q);
 
    memset(chk, 0, sizeof(chk));
    memset(cnt, 0, sizeof(cnt));
    V.clear();
    for(i=1; i<=N; i++) if(DP[i]>DQ[i]) chk[i]=1;
    dfs2(P, P);
    dfs3(P, P, 0);
 
    memset(chk, 0, sizeof(chk));
    memset(cnt, 0, sizeof(cnt));
    V.clear();
    for(i=1; i<=N; i++) if(DP[i]<=DQ[i]) chk[i]=1;
    dfs2(Q, Q);
    dfs3(Q, Q, 0);    
 
    for(i=1; i<=N; i++) printf("%d\n", ans[i]);
}