[SWEA] 1855 영준이의 진짜 BFS

 

문제를 잘 읽어보면 한가지 알고리즘이 떠오른다

LCA!!! 간만에 떠올라서 잠깐 다시 공부하고 오랜만에 만들어봤다!

 

먼저 이 문제는 1이 루트고 bfs로 나머지 정점들에 순서대로 계속 가기 때문에

정점 2개씩 잡으며 lca를 구해서 lca까지의 거리를 더해가면 답을 구할 수 있다

 

lca만 알면 쉽게 풀 수 있는 문제다!

 

근데 더 어려운건 swea에서 스택메모리가 1메가라그런지 dfs로 깊이를 매기니

런타임 에러가 났다;; 대체 어디가 틀린지 몰랐는데 이게 틀렸었다;;

 

bfs로 바꿔서 깊이를 찾았다!!

 

함수였던건 다 메인으로 넣어버렸다!

 

소스코드

#include <iostream>
#include <algorithm>
#include <vector>
#include <cstring>
#include <queue>
using namespace std;
#define ll long long
const int MAXN = 100005;
int n;
int d[MAXN],p[MAXN][23];
bool check[MAXN];
vector<vector<int>> Graph;
 
void sparse(){
    for(int i=1;i<=20;i++){
        for(int j=0;j<n;j++){
            p[j][i] = p[p[j][i-1]][i-1];
        }
    }
}
 
int lca(int x,int y){
    if(d[y]>d[x]) swap(x,y);
    for(int i=19;i>=0;i--){
        if((d[x]-d[y])>=(1<<i)) x = p[x][i];
    }
    if(x==y) return x;
    for(int i=19;i>=0;i--){
        if(p[x][i] != p[y][i]){
            x = p[x][i];
            y = p[y][i];
        }
    }
    return p[x][0];
}
 
int main(){
    //freopen("input.txt","r",stdin);
    int tcase; scanf(" %d",&tcase); int c=1;
    while(tcase--){
        memset(check,0,sizeof(check));
        memset(d,0,sizeof(d));
        memset(p,0,sizeof(p));
        Graph.clear();
 
        scanf(" %d",&n);
        Graph.resize(n+3);
 
        for(int i=1;i<n;i++){
            int x; scanf(" %d",&x);
            Graph[x-1].push_back(i);
        }
        for(int i=0;i<n;i++) sort(Graph[i].begin(),Graph[i].end());
 
        queue<pair<int,int>> dq;
        dq.push({0,0});
        check[0]=true;
 
        while(!dq.empty()){
            int now = dq.front().first;
            int depth = dq.front().second;
            dq.pop();
            d[now]=depth;
 
            for(int i=0;i<Graph[now].size();i++){
                int next = Graph[now][i];
                if(!check[next]){
                    check[next]=true;
                    p[next][0]=now;
                    dq.push({next,depth+1});
                }
            }
        }
       
        sparse();
 
        ll ret=0;
        memset(check,0,sizeof(check));
        vector<int> vt;
        queue<int> q;
        q.push(0);
        check[0]=true;
        
        while(!q.empty()){
            int now = q.front();
            q.pop();
 
            vt.push_back(now);
            for(int i=0;i<Graph[now].size();i++){
                int next = Graph[now][i];
                if(!check[next]){
                    check[next]=true;
                    q.push(next);
                }
            }
        }
        for(int i=0;i<n-1;i++){
            int LCA = lca(vt[i],vt[i+1]);
            ll a = d[vt[i]] - d[LCA];
            ll b = d[vt[i+1]] - d[LCA];
            ret+=a; ret+=b;
        }
 
        printf("#%d %lld\n",c++,ret);
    }
}