백준 14942번: 개미 [C++]

작성일 | In
플래티넘 V 플래티넘 V

문제

개미
희소 배열과 이분 탐색으로 해결할 수 있는 문제입니다.

풀이

굴을 타고 한 방에서 다른 방으로 갈 수 있는 경로가 유일하게 존재한다고 합니다. 따라서 개미굴이 트리 구조임을 알 수 있습니다.
만약 개미굴이 사이클을 포함한다면 그 사이클 내에서 두 방 사이의 경로는 최소 2개 이상 존재하기 때문입니다.

그렇다면 희소 배열과 이분탐색을 활용하여 문제를 풀 수 있습니다.

코드

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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
#include <bits/stdc++.h>
using namespace std;


vector<int> energes; // 개미들의 에너지
vector<vector<pair<int, int>>> adj // 인접 리스트
vector<vector<int>> parent; // parent[node][i]는 node에서 2^i만큼 부모로 올라갔을 때 노드
vector<vector<int>> dist; // dist[node][i]는 node에서 2^i만큼 떨어진 부모와의 거리(length)
vector<int> depth; // 노드의 깊이
vector<int> visited;

void setupDist(int now) {
    for (auto [next, cost] : adj[now]) {
        if (visited[next]) continue;
        visited[next] = true;
        
        parent[next][0] = now;
        dist[next][0] = cost;
        depth[next] = depth[now] + 1;

        for (int i = 1; i < 18; i++) {
            if (parent[next][i-1] == 0) break;
            parent[next][i] = parent[parent[next][i-1]][i-1];
            dist[next][i] = dist[next][i-1] + dist[parent[next][i-1]][i-1];
        }

        setupDist(next);
    }
}

int binarySearch(int x) {
    int lo = 0, hi = depth[x] + 1; // 몇 번까지 움직일 수 있는지 [lo, hi)

    while (lo < hi) { // upper_bound로 탐색
        int mid = (lo + hi) / 2;
        int length = 0;

        int now = x;
        for (int i = 0; i < 18; i++) {
            if (mid & (1 << i)) {
                length += dist[now][i];
                now = parent[now][i];
            }
        }

        if (length <= energes[x]) lo = mid + 1;
        else hi = mid;
    }
    lo--; // upper_bound니까 1 빼기
    
    int ret = x;
    for (int i = 0; i < 18; i++) {
        if (lo & (1 << i)) ret = parent[ret][i];
    }
    return ret;
}

int main() {
    // NOTE: 1번 대신 0번 사용
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n;
    cin >> n;
    energes.resize(n);
    for (int i = 0; i < n; i++) cin >> energes[i];

    adj.assign(n, vector<pair<int, int>>());
    for (int i = 0; i < n-1; i++) {
        int a, b, c;
        cin >> a >> b >> c;
        a--, b--;
        adj[a].emplace_back(b, c);
        adj[b].emplace_back(a, c);
    }
    parent.resize(n, vector<int>(18)); // 2^17 = 131072이기 때문 
    dist.resize(n, vector<int>(18));
    depth.resize(n);
    visited.resize(n);

    visited[0] = true;
    setupDist(0);
    for (int i = 0; i < n; i++) {
        cout << binarySearch(i) + 1 << '\n';
    }
}

백준 14942번: 개미 [C++]