[백준 20297] Confuzzle

작성일 | In |

문제 링크

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

풀이

센트로이드의 아이디어를 조금 가져와 보죠. 해당 정점을 지나는 경로중 정답과 가장 가까운 거리가 가장 짧은 값을 관리합니다. dfs를 하면서도 해당 정점을 지나는 경로를 관리할 수는 없을까요? 정점에서 “자식들에 저장된 수와 그 수가 저장된 정점 중 가장 낮은 (제 코드에서는 root의 높이가 1입니다.) 높이”를 관리한다고 합시다. 이 값들을 가지고 있다면 현재 정점에서 자식 정점끼리 합쳐서 현재 정점을 지나는 가장 짧은 경로를 뽑아낼 수 있습니다. 그러므로 이 값을 관리하고 합치기만 잘 한다면 정답을 뽑아낼 수 있습니다.

수는 최대 10만개 까지 있을 수 있으므로 배열로 관리하기에는 어려울 겁니다. (공간복잡도 $O(N^2)$, 합치는 데는 더욱 큰 시간복잡도가 소요됩니다.) 그러므로 저장하는 수만큼의 메모리를 사용하는 map을 이용해서 값을 저장할 것입니다. 메모리의 문제는 해결했지만 여전히 자식노드들을 합치는 데에는 큰 시간이 필요합니다.

small to large 테크닉을 이용해서 map을 합치는 데에 총 시간복잡도는 $O(N\log^2N)$이 걸리고 값을 합치면서 해당 정점을 포함하는 가장 짧은 경로를 뽑아낼 수 있으므로 문제의 답 또한 알아낼 수 있습니다. 자세한 구현은 코드를 참고해주세요.

전체 코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
#include<bits/stdc++.h>
using namespace std;
using ll = long long;

int n, c[100100], id[100100], lev[100100], mn = 1e9;
vector<int> v[100100];
map<int, int> mp[100100];

void dfs(int u, int p){
    lev[u] = lev[p] + 1;

    if(p && v[u].size() == 1){
        id[u] = u;
        mp[id[u]][c[u]] = lev[u];
        return;
    }

    int mx = 0;
    for(auto x : v[u]) if(x != p){
        dfs(x, u);
        if(mp[id[x]].size() > mx) mx = mp[id[x]].size(), id[u] = id[x];
    }

    if(mp[id[u]].count(c[u])) mn = min(mn, mp[id[u]][c[u]] - lev[u]);
    mp[id[u]][c[u]] = lev[u];

    for(auto x : v[u]) if(x != p && id[x] != id[u]){
        for(auto [y, z] : mp[id[x]]){
            if(mp[id[u]].count(y)){
                mn = min(mn, mp[id[u]][y] + z - lev[u] * 2);
                mp[id[u]][y] = min(mp[id[u]][y], z);
            }else{
                mp[id[u]][y] = z;
            }
        }
    }
}

int main(){
    scanf("%d",&n);
    for(int i=1;i<=n;i++) scanf("%d",c+i);
    for(int i=1;i<n;i++){
        int a, b; scanf("%d %d",&a,&b);
        v[a].push_back(b), v[b].push_back(a);
    }
    dfs(1, 0);
    printf("%d",mn);
}

여담

문제를 보고 바로 센트로이드를 이용한 풀이가 떠올랐지만 이 문제에서는 업데이트가 없기 때문에 다른 방식으로 접근이 가능할 거라 생각하고 이 풀이를 떠올렸습니다.