백준 20297 Confuzzle

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

 

먼가 문제 이름이 익숙한 걸 보니 전에도 몇 번 보고 넘겼었던 것 같다. 문제를 읽고 가장 먼저 트리 압축 풀이가 떠올랐고, 그 다음으로는 센트로이드 풀이가 떠올랐고, 업데이트 없는 센트로이드질이니 small to large 풀이가 떠올랐다. 루트질로 풀 수 있다고 해서 더 고민해봤다. 해당 테크닉을 안다면 다른 세 가지 방법 중 하나로 푸는 게 편할 것 같다. 사실 네 가지 풀이를 다 짜 보고 블로그에 정리하고 싶었는데 그런 글이 이미 있더라.. 

 

값 c가 같은 정점들끼리 그룹을 짓고, 각 그룹에 대해 독립적으로 문제를 푼다고 생각하자. 어떤 그룹에 속한 정점 중 가장 가까운 쌍을 찾는 방법 두 가지를 생각할 수 있다.

1. 그룹에서 정점 2개를 고르는 모든 경우의 수마다 거리를 계산한다. 
2. dfs를 수행하며 자손 중 해당 그룹에 속한 가장 가까운 정점 두 개를 관리하는 트리 DP를 수행한다.

트리의 전체 정점이 $N$개, 해당 그룹의 정점 수를 $X$라고 하자. 1번 방법은 $O(X^2\log N)$, 2번 방법은 $O(N)$이다. 임의의 상수 $B$를 잡아 그룹의 정점 수가 $B$ 이하라면 1번, 초과라면 2번 방법을 사용하자. 1번 방법을 사용할 그룹은 최대 $O(N)$개이므로 총 시간복잡도는 $O(NB^2\log N)$, 2번 방법을 사용할 그룹은 최대 $\frac{N}{B}$개이므로 총 시간복잡도는 $O(\frac{N^2}{B})$이다. $B = \sqrt{N}$으로 잡으면 전체 시간복잡도가 $O(N\sqrt{N} \log N)$이 되어 문제를 풀 수 있다. 물론 lca를 $O(1)$로 구현하면 로그도 뗄 수 있다.

 

아래 코드는 틀렸습니다! 를 받은 코드로, 디버깅에 거의 하루가 걸렸다. 심심한 독자분들은 어떤 실수가 있는지 찾아보시기 바랍니다 

#include <bits/stdc++.h>
using namespace std;
 
int c[101010], par[18][101010], d[101010], N, ans = 1e9;
const int B = 333;
vector<int> T[101010], V[101010];
 
int lca(int a, int b){
    if (d[a] > d[b]) swap(a, b);
    for (int i = 0; d[a] < d[b]; i++){
        if ((d[b] - d[a]) & (1 << i))
            b = par[i][b];
    }
    if (a == b) return a;
    for (int i = 17; i >= 0; i--){
        if (par[i][a] != par[i][b])
            a = par[i][a], b = par[i][b];
    }
    return par[0][a];
}
 
int dfs(int v, int now){
    vector<int> dist{(int)1e9,(int)1e9};
    for (int i : T[v]){
        if (i == par[0][v]) continue;
        d[i] = d[v] + 1;
        par[0][i] = v;
        dist.push_back(dfs(i, now) + 1);
    }
    sort(dist.begin(), dist.end());
    if (c[v] == now) {
        ans = min(ans, dist[0]);
        return 0;
    }
    else {
        ans = min(ans, dist[0]+dist[1]);
        return dist[0];
    }
}
 
int main() {
    cin.tie(nullptr)->sync_with_stdio(0);
 
    cin >> N;
    for (int i = 1; i <= N; i++){
        cin >> c[i];
        V[c[i]].push_back(i);
    }
    for (int i = 1,u,v; i < N; i++){
        cin >> u >> v;
        T[u].push_back(v);
        T[v].push_back(u);
    }
    dfs(1, 1);
    for (int i = 2; i <= N; i++){
        if (V[i].size() <= B) continue;
        dfs(1, i);
    }
    for (int i = 1; i <= 17; i++)
        for (int j = 1; j <= N; j++)
            par[i][j] = par[i-1][par[i-1][j]];
    for (int i = 1; i <= N; i++){
        if (V[i].size() <= B){
            for (int a = 0; a < V[i].size(); a++)
                for (int b = a+1; b < V[i].size(); b++)
                    ans = min(ans, d[V[i][a]] + d[V[i][b]] - d[lca(V[i][a],V[i][b])]);
        }
    }
    cout << ans;
}