[C++] 13511번 트리와 쿼리 2 - 자료 구조, 트리, 최소 공통 조상, 희소 배열

1. 문제

• N개의 중점으로 된 트리가 있다. 1~N번 매겨져 있고 간선은 1~N-1번 매겨져 있다. 아래 두 쿼리를 수행하는 프로그램을 작성하시오.
  • 1 u v : u->v 경로 비용 출력
  • 2 u v k : u->v 경로 정점 중 k번째 정점 출력
    • k는 u->v 경로에 포함된 정점 수보다 작다.

<입력>

• - 1 -   \(N(2 \leq N \leq 100,000)\)
• - N-1 개줄 -   i번 간선 u->v와 비용 \(w(1 \leq w \leq 1,000,000)\)
• - N+1 -   쿼리의 개수 \(M(1 \leq M \leq 100,000)\)
• - M개의 줄 -   쿼리

<출력>

• 각 쿼리 결과 출력

2. 재정의

• X

3. 해결 방법

• 트리 노드 차수 degree 설정
• 희소 배열 비용, 노드 메모
• 비용 : LCA 하면서 더하기
• 노드 : k--. degree[u] - degree[LCAnode] <= k이면 v를 기준으로 k = degree[u] + degree[v] - 2*degree[LCAnode] - k. u=v 해서 희소 배열 써서 구하면 됨

4. 실수한 점, 개선할 점

• k값 구하는거 그냥 때려박을려 하지 말고 생각해서 직접 구해서 넣는 게 더 시간 절약 됐을 것 같다.
• 알고리즘이 똑같은데 40ms 정도 차이난다. 아마 LCA를 반환하면서 생긴 시간 차이인 것 같다. 더 시간 소모가 짧은 코드는 아래에 링크해두었다.

 

<코드>

#include <iostream>
#include <cmath>
#include <vector>

#define pil pair<int, long long>
#define pii pair<int, int>
#define fi first
#define se second

using namespace std;
typedef long long ll;

const int MAX_N = 100001;
const int MAX_Sparse_Index = 17;

// <문제>
// 정점의 개수 N, 쿼리의 개수 M
int N, M;
// 가중치가 있는 트리
vector<pii> adj[MAX_N];

// LCA
int degree[MAX_N];
// Sparse Array
pil parent[MAX_N][MAX_Sparse_Index];
int SparseIndexMax;

// 희소 배열 전 작업. 부모 노드와 노드 차수 구하기
void dfs(int no) {
    for(auto next : adj[no]) {
        // 탐색하지 않은 노드인 경우
        if(!degree[next.fi]) {
            parent[next.fi][0].fi = no;
            parent[next.fi][0].se = next.se;
            degree[next.fi] = degree[no] + 1;
            dfs(next.fi);
        }
    }
}

void input() {
    cin >> N;
    SparseIndexMax = (int)ceil(log2(N));
    
    // 무방향 가중치 그래프 그리기
    for(int i=0; i<N-1; i++) {
        int u, v, w;
        cin >> u >> v >> w;
        adj[u].push_back({v, w});
        adj[v].push_back({u, w});
    }
    
    // 희소 배열 전 작업. 부모 노드와 노드 차수 구하기
    degree[1] = 1;
    dfs(1);
    
    // 희소 배열
    for(int i=1; i<SparseIndexMax; i++) {
        for(int j=1; j<=N; j++) {
            parent[j][i].fi = parent[parent[j][i-1].fi][i-1].fi;
            parent[j][i].se = parent[j][i-1].se + parent[parent[j][i-1].fi][i-1].se;
        }
    }
}

// 공통 조상 찾기
pil LCA(int u, int v) {
    ll ret = 0;
    
    if(degree[u] < degree[v])
        swap(u, v);
    
    int diff = degree[u] - degree[v];
    for(int i=0; diff; i++) {
        if(diff%2) {
            ret += parent[u][i].se;
            u = parent[u][i].fi;
        }
        diff >>= 1;
    }
    
    if(u != v) {
        for(int i=SparseIndexMax - 1; i>=0; i--) {
            if(parent[u][i].fi && parent[u][i].fi != parent[v][i].fi) {
                ret += parent[u][i].se + parent[v][i].se;
                u = parent[u][i].fi;
                v = parent[v][i].fi;
            }
        }
        ret += parent[u][0].se + parent[v][0].se;
        u = parent[u][0].fi;
    }
    
    return {u, ret};
}

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);
    
    input();
    
    cin >> M;
    while(M--) {
        int what;
        cin >> what;
        // 1번 쿼리
        if(what & 1) {
            int u, v;
            cin >> u >> v;
            // u->v 비용 출력
            cout << LCA(u, v).se << '\n';
        }
        // 2번 쿼리
        else {
            int u, v, k;
            cin >> u >> v >> k;
            
            k--;
            int LCANode = LCA(u, v).fi;
            if(degree[u] - degree[LCANode] <= k) {
                k = degree[u] + degree[v] - 2*degree[LCANode] - k;
                u = v;
            }
            
            for(int i=0; k>0; i++) {
                if(k%2)
                    u = parent[u][i].fi;
                
                k >>= 1;
            }
            cout << u << '\n';
        }
    }
    
    return 0;
}

 

<LCA 알고리즘>

https://kibbomi.tistory.com/201

최소 공통 조상 (LCA, Lowest Common Ancestor)(C/C++)

 

<C++ ceil()>

https://blockdmask.tistory.com/112

[C언어/C++] 올림, 내림, 반올림 (floor, ceil) 함수

 

<참고한 코드>

https://www.acmicpc.net/source/52295621

로그인

 

<문제 바로가기>

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

13511번: 트리와 쿼리 2

 

※현재 고등학교 등교 중인 학생입니다. 이제 알고리즘을 본격적으로 공부하기 시작해서 아직 초보입니다. 혹시 제가 잘못 알고 있는 점이나 더 좋은 풀이 방법이 있어 댓글에 남겨주시면 감사히 하나하나 열심히 읽어보겠습니다. 좋아요, 단순한 댓글 한마디라도 저에겐 큰 힘이 됩니다! 감사합니다.