백준 2927번: 남극 탐험 [C++]

작성일 | In
다이아몬드 V 다이아몬드 V

문제

남극 탐험

풀이

이전까지 지어진 다리를 통해 이동할 수 있는지 판단하는 것(bridge 쿼리)은 분리 집합을 통해 할 수 있고, penguins 쿼리와 excursion 쿼리는 heavy-light 분할 + 세그먼트 트리로 해결할 수 있습니다.

먼저 쿼리를 모두 입력받아 트리를 구성한 후 쿼리를 처리해주면 됩니다.

주어지는 입력이 트리 하나가 아니라 포레스트일 수 있음에 주의합시다.

코드

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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
#include <bits/stdc++.h>
#ifndef ONLINE_JUDGE
#define ASSERT(x) assert(x)
#else
#define ASSERT(ignore) ((void)0)
#endif

using namespace std;

// Segment Tree
struct SegmentTree {
    int n;
    vector<int> t;

    SegmentTree(int size) : n(size), t(size << 1) {}

    void modify(int i, int v) {
        for (t[i += n] = v; i > 1; i >>= 1)
            t[i >> 1] = t[i] + t[i ^ 1];
    }

    int getRangeSum(int l, int r) {
        int ret = 0;

        for (l += n, r += n; l <= r; l >>= 1, r >>= 1) {
            if (l & 1)
                ret += t[l++];
            if (~r & 1)
                ret += t[r--];
        }

        return ret;
    }
};

// Disjoint Set
struct DisjointSet {
    vector<int> parent;

    DisjointSet(int size) : parent(size) {
        for (int i = 0; i < size; i++)
            parent[i] = i;
    }

    int findRoot(int x) {
        if (x != parent[x])
            parent[x] = findRoot(parent[x]);
        return parent[x];
    }

    void unionRoot(int u, int v) { parent[findRoot(u)] = findRoot(v); }

    bool connected(int u, int v) { return findRoot(u) == findRoot(v); }
};

// Heavy-light Decomposition
vector<int> parent, sz, depth, in, top, penguinCnt;
vector<vector<int>> adj;
vector<bool> initialized;

void initHLD(int node) {
    initialized[node] = true;

    for (int& child : adj[node]) {
        if (child != parent[node])
            continue;

        swap(child, adj[node].back());
        adj[node].pop_back();
        break;
    }

    sz[node] = 1;
    for (int& child : adj[node]) {
        depth[child] = depth[node] + 1;
        parent[child] = node;
        initHLD(child);
        sz[node] += sz[child];

        if (sz[child] > sz[adj[node].front()])
            swap(child, adj[node].front());
    }
}

void initTour(int node) {
    static int visitedCnt = 0;
    in[node] = ++visitedCnt;

    for (int child : adj[node]) {
        top[child] = (child == adj[node].front() ? top[node] : child);
        initTour(child);
    }
}

// Queries
enum QueryType : uint8_t {
    BRIDGE,
    PENGUINS,
    EXCURSION,
};
struct Query {
    QueryType type;
    int args[2];
};
vector<Query> queries;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int N;
    cin >> N;
    parent = sz = depth = in = top = penguinCnt = vector<int>(N + 1);
    adj.resize(N + 1);
    initialized.resize(N + 1);
    for (int i = 1; i <= N; i++)
        cin >> penguinCnt[i];

    int Q;
    cin >> Q;
    queries.resize(Q);
    DisjointSet disjointSet(N + 1);
    for (Query& q : queries) {
        string type;
        cin >> type >> q.args[0] >> q.args[1];

        if (type == "bridge") {
            q.type = BRIDGE;
            auto [A, B] = q.args;
            if (not disjointSet.connected(A, B)) {
                adj[A].push_back(B);
                adj[B].push_back(A);
                disjointSet.unionRoot(A, B);
            }
        } else if (type == "penguins")
            q.type = PENGUINS;
        else /* if (type == "excursion") */
            q.type = EXCURSION;
    }

    for (int i = 1; i <= N; i++) {
        if (not initialized[i]) {
            initHLD(i);
            top[i] = i;
            initTour(i);
        }
    }

    SegmentTree segTree(N + 1);
    for (int i = 1; i <= N; i++) {
        segTree.modify(in[i], penguinCnt[i]);
    }

    disjointSet = DisjointSet(N + 1);
    for (Query q : queries) {
        switch (q.type) {
            case BRIDGE: {
                auto [A, B] = q.args;
                println("{}", disjointSet.connected(A, B) ? "no" : "yes");
                disjointSet.unionRoot(A, B);
                break;
            }
            case PENGUINS: {
                auto [A, X] = q.args;
                segTree.modify(in[A], X);
                break;
            }
            case EXCURSION: {
                auto [A, B] = q.args;
                if (not disjointSet.connected(A, B)) {
                    println("impossible");
                    break;
                }

                int ans = 0;
                while (top[A] != top[B]) {
                    if (depth[top[A]] > depth[top[B]])
                        swap(A, B);
                    ans += segTree.getRangeSum(in[top[B]], in[B]);
                    B = parent[top[B]];
                }
                if (depth[A] > depth[B])
                    swap(A, B);
                ans += segTree.getRangeSum(in[A], in[B]);
                println("{}", ans);
                break;
            };
        }
    }

    return 0;
}

백준 2927번: 남극 탐험 [C++]