백준21815 Through Another Maze Darkly

작성일 | In

문제 링크

  • http://icpc.me/21825

문제 출처

  • 2021 CCO Day1 3번

사용 알고리즘

  • DFS

시간복잡도

  • $$

풀이

다음과 같은 트리를 생각해 봅시다. 각 정점에서 가장 먼저 사용하는 간선을 화살표로 표현했습니다.

첫 번째 순회에서는 1 2 5 2 6 8 6 2 1 3 1 순서대로 방문합니다. 첫 번째 순회를 한 이후의 트리는 다음과 같습니다. 첫 번째 순회에서 방문한 정점들의 화살표가 모두 첫 번째 자식(없다면 부모) 정점을 가리키고 있다는 점을 관찰할 수 있습니다.

두 번째 순회에서는 1 2 4 2 5 2 6 8 6 2 1 3 7 3 1 순서대로 방문합니다. $t-1$번째 순회에서 방문한 정점은 항상 $t$번째 순회에서 방문한다는 점도 알 수 있습니다. 또한, 각 순회에서 정점을 방문하는 순서가 역전되지 않습니다.
따라서 자식 정점들을 적당한 순서대로 배치한 다음 오일러 투어를 구해 놓으면, 두 정수 $t, k$가 주어졌을 때 $t$번째 순회에서 $k$번째로 방문하는 정점을 세그먼트 트리 등을 이용해 오프라인으로 구할 수 있습니다.

입력으로 트리의 인접 리스트가 주어지면 부모 정점의 위치 $p$를 구합시다. $p$ 이전에 있는 정점은 현재 정점을 처음으로 방문한 순회에서, $p$ 이후에 있는 정점은 그 다음 순회에서 방문하게 됩니다. 따라서 DFS를 이용해 $O(N)$ 시간에 각 정점을 처음으로 방문하는 시점을 구할 수 있습니다. 각 간선을 처음 사용하게 되는 시점도 함께 구할 수 있습니다.
이후 구현의 편의를 위해 부모 정점이 가장 마지막에 오도록 인접 리스트를 cyclic shift 합시다.

이제 $K$번째로 방문하는 정점을 구하는 쿼리를 처리할 차례입니다. 우리의 궁극적인 목표는 $K$를 적당한 정수 $t, k$로 바꿔서, $t$번째 순회에서 $k$번째로 방문하는 정점을 구하는 문제로 변환하는 것입니다.
각 간선이 처음 사용되는 시점을 알고 있다면, 누적 합과 비슷한 방식으로 각 순회의 길이를 상수 시간에 구할 수 있습니다. 따라서 순회의 길이의 누적합 위에서 이분 탐색을 하면 $O(\log N)$ 시간에 $t$를 구할 수 있습니다. $N+1$번째 순회부터는 항상 모든 정점과 간선을 방문한다는 것에 유의하세요.
$t$를 구했다면 $k$를 구하는 것은 간단합니다. 단순히 모듈러 연산을 이용하면 되기 때문입니다.

$t, k$가 주어졌을 때 실제 정답을 찾는 것은 쿼리를 $t$ 오름차순으로 정렬한 다음, 세그먼트 트리에서 $k$번째 원소를 찾는 연산을 사용하면 됩니다. 자세한 구현은 코드의 64번째 줄부터 참고하시면 됩니다.

전체 코드

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
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
constexpr int SZ = 1 << 21;

int N, Q, On[808080], R[808080], T[SZ];
vector<int> G[808080];
vector<pair<int,int>> Tour;
vector<int> Idx[808080]; // Idx[t] : t에 켜지는 정점
ll TourSize[808080];

void Add(int x, int v){ for(x+=3; x<SZ; x+=x&-x) T[x] += v; }
int Get(int x){ int ret = 0; for(x+=3; x; x-=x&-x) ret += T[x]; return ret; }
int Kth(int x){
    int l = 0, r = Tour.size() - 1;
    while(l < r){
        int m = (l + r) / 2;
        if(Get(m) >= x) r = m;
        else l = m + 1;
    }
    return r;
}

void GetTime(int v, int b=-1, int t=1){
    On[v] = t;
    int pos = find(G[v].begin(), G[v].end(), b) - G[v].begin();
    for(int i=0; i<pos; i++) GetTime(G[v][i], v, t);
    for(int i=pos+1; i<G[v].size(); i++) GetTime(G[v][i], v, t+1);
    if(b != -1) rotate(G[v].begin(), G[v].begin()+pos, G[v].end());
}

inline void AddEdge(int s, int e){
    Idx[max(On[s], On[e])].push_back(Tour.size());
    Tour.emplace_back(s, e);
}

void GetTour(int v, int b=-1){
    for(auto i : G[v]) if(i != b) AddEdge(v, i), GetTour(i, v), AddEdge(i, v);
}

int main(){
    ios_base::sync_with_stdio(false); cin.tie(nullptr);
    cin >> N >> Q;
    for(int i=1; i<=N; i++){
        int sz; cin >> sz; G[i].resize(sz);
        for(auto &j : G[i]) cin >> j;
        rotate(G[i].begin(), G[i].begin()+1, G[i].end());
    }
    GetTime(1);
    GetTour(1);
    for(int i=1; i<=N; i++) TourSize[i] = TourSize[i-1] + Idx[i].size();
    for(int i=1; i<=N; i++) TourSize[i] += TourSize[i-1];

    vector<tuple<ll,ll,ll>> V; // iteration, kth, idx
    for(int q=1; q<=Q; q++){
        ll K; cin >> K;
        int it = lower_bound(TourSize+1, TourSize+N+1, K) - TourSize;
        K -= TourSize[it-1];
        if(it == N+1) K = (K - 1) % Tour.size() + 1;
        V.emplace_back(it, K, q);
    }
    sort(V.begin(), V.end());

    for(int it=1, i=0; it<=N+1; it++){
        for(auto e : Idx[it]) Add(e, 1);
        while(i < V.size() && get<0>(V[i]) == it){
            R[get<2>(V[i])] = Kth(get<1>(V[i]));
            i += 1;
        }
    }
    for(int i=1; i<=Q; i++) cout << Tour[R[i]].second << " ";
}