백준 24520 - Meet In The Middle (C++)

문제

• 문제 링크

  
   BOJ 24520 - Meet In The Middle 
  

• 문제 요약

$N$개의 마을이 트리 형태로 주어지고 $Q$개의 약속이 마을 쌍의 형태로 주어진다.
각 약속마다, 두 마을에서 정확히 같은 거리만큼 떨어진 마을의 위치를 구해보자.

• 제한

TL : $1$ sec, ML : $1024$ MB




$1 ≤ N, Q ≤ 10^5$

$1 ≤ w ≤ 10^9$

알고리즘 분류

• 자료 구조 (data_structures)
• 트리 (trees)
• 최소 공통 조상 (lowest common ancestor)
• 희소 배열 (sparse_table)

풀이

$dist[x]$ : 루트로부터 $x$까지의 거리.

$dp[x][y]$ : $x$의 $2^y$번째 부모 정점.

라고 할 때 $O(NlogN)$의 전처리를 통해 어떤 두 정점의 $LCA$를 $O(logN)$만에 구할 수 있다.
나아가 두 정점 사이의 거리도 아래의 식으로 $O(logN)$에 구할 수 있다.

$dist_{xy} = dist[x] + dist[y] - 2 * dist[lca(x, y)]$

---

이제 쿼리마다 두 정점의 거리는 쉽게 구할 수 있는데, 두 정점으로부터 정확히 같은 거리만큼 떨어진 정점은 어떻게 찾을 수 있을까$?$

이 과정에서도 $LCA$를 찾을 때와 비슷하게 $\frac{dist_{xy}}{2}$ 만큼 최대한 올리는 $Binary$ $Lifting$을 적용해볼 수 있다.

• $LCA$를 찾을 때처럼 일반성을 유지한 채 루트로부터 더 멀리 있는 정점을 $x$, 나머지를 $y$라 하자.
• 시작 정점을 $cur == x$로 두고 $dp[cur][i]$에 기록된 가장 높은 정점의 정보로부터 차례로 내려오며 아래의 조건을 검사하자.
• $dp[cur][i]$가 $x$에서 $lca_{xy}$에 이르는 경로 상에 위치한지$?$
• 그만큼 뛰어도 되는지 즉, $dist[x] - dist[dp[cur][i]] ≤ \frac{dist_{xy}}{2}$ 를 만족하는지$?$

그렇게 최종적으로 위치를 잡은 $cur$에 대해 $dist[x] - dist[cur]$ 의 값을 보고 답을 결정해주면 된다.

아 추가로 위의 과정은 $dist_{xy}$가 짝수일 때만 해당된다.
거리가 홀수라면 중간 정점을 찾을 수 없음이 자명하니 말이다.

전체 코드

#include<bits/stdc++.h>
#define N 100001
using namespace std;
 
vector <pair<int, int>> Gr[N];
int dep[N], dp[N][18];
long long dist[N];
int n, q;
 
void Input()
{
    ios_base::sync_with_stdio(0); cin.tie(0);
    cin >> n >> q;
    for (int x, y, w, i{}; ++i < n;)
    {
        cin >> x >> y >> w;
        Gr[x].emplace_back(y, w);
        Gr[y].emplace_back(x, w);
    }
}
void LCAInit(int now, int depth)
{
    dep[now] = depth;
    for (auto& [next, weight] : Gr[now])
        if (!dep[next])
        {
            dist[next] = dist[now] + weight;
            dp[next][0] = now;
            LCAInit(next, depth + 1);
        }
}
void TableDP()
{
    for (int j(1); j < 18; j++)
        for (int i(1); i <= n; i++)
            dp[i][j] = dp[dp[i][j - 1]][j - 1];
}
int getLCA(int x, int y)
{
    if (dep[x] < dep[y]) swap(x, y);
    int diff(dep[x] - dep[y]);
 
    for (int i{}; diff; diff >>= 1, i++)
        if (diff & 1)
            x = dp[x][i];
    if (x ^ y)
    {
        for (int i(17); !!~i; i--)
            if (dp[x][i] ^ dp[y][i])
                x = dp[x][i], y = dp[y][i];
        x = dp[x][0];
    }
    return x;
}
void Query()
{
    for (int x, y; q--;)
    {
        cin >> x >> y;
        int lca(getLCA(x, y));
 
        long long dis(dist[x] + dist[y] - 2 * dist[lca]);
 
        if (dis & 1) cout << -1 << '\n';
        else
        {
            if (dist[x] - dist[lca] < dist[y] - dist[lca]) swap(x, y);
            int cur(x); dis >>= 1;
 
            for (int i(17); !!~i; i--)
                if (dep[dp[cur][i]] >= dep[lca] && dist[x] - dist[dp[cur][i]] <= dis)
                    cur = dp[cur][i];
 
            cout << (dist[x] - dist[cur] == dis ? cur : -1) << '\n';
        }
    }
}
int main()
{
    Input();
    LCAInit(1, 1);
    TableDP();
    Query();
}

comment

끝