중급반 7주차 풀이 - DSU, MST
1. [1717] 집합의 표현 (Gold IV)
Union-Find 구현 문제입니다. 특별한 것도 없습니다.
int par[1000020];
int fnd(int a){
if (par[a] == a){ return a; }
return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }
void Main(){
int n, q; cin >> n >> q;
for (int i = 0; i <= n; i++){ par[i] = i; }
while (q--){
int typ; cin >> typ;
if (typ == 0){
int a, b; cin >> a >> b; uni(a, b);
}
if (typ == 1){
int a, b; cin >> a >> b;
cout << ( fnd(a) == fnd(b) ? "YES" : "NO" ) << endl;
}
}
}2. [1976] 여행 가자 (Gold IV)
연결된 도시들끼리 모두 Union을 해봅시다.
그럼 같은 집합에 속한 도시들끼리는 서로 여행을 갈 수 있고, 그렇지 않으면 서로 여행을 갈 수 없게 됩니다.
그러니 여행 계획에서 다른 집합으로 넘어가야 하는 일이 발생하는지 체크해주면 됩니다.
int par[220];
int fnd(int a){
if (par[a] == a){ return a; }
return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }
int arr[1020];
void Main(){
int n, m; cin >> n >> m;
for (int i = 1; i <= n; i++){ par[i] = i; }
for (int i = 1; i <= n; i++){
for (int j = 1; j <= n; j++){
int x; cin >> x; if (x){ uni(i, j); }
}
}
for (int i = 1; i <= m; i++){ cin >> arr[i]; }
for (int i = 1; i < m; i++){
if ( fnd(arr[i]) != fnd(arr[i+1]) ){ cout << "NO"; return; }
}
cout << "YES"; return;
}3. [18116] 로봇 조립 (Gold IV)
이번에는 집합의 크기를 추가적으로 저장해야 합니다.
초기 상태에서는 모든 집합의 크기가 1이고, 두 다른 집합을 합치면 새로운 집합의 크기는 두 집합 크기의 합이 되므로 이를 그대로 구현해주면 됩니다.
const int N = 1000000;
int par[1000020], siz[1000020];
int fnd(int a){
if (par[a] == a){ return a; }
return par[a] = fnd(par[a]);
}
void uni(int a, int b){ int pa = fnd(a), pb = fnd(b);
if (pa == pb){ return; }
par[pa] = pb; siz[pb] += siz[pa];
}
void Main(){
for (int i = 1; i <= N; i++){ par[i] = i; siz[i] = 1; }
int q; cin >> q; while (q--){
char typ; cin >> typ;
if (typ == 'I'){ int a, b; cin >> a >> b; uni(a, b); }
if (typ == 'Q'){ int a; cin >> a; cout << siz[fnd(a)] << endl; }
}
}4. [1197] 최소 스패닝 트리 (Gold IV)
MST 연습 문제입니다. Kruskal로 풀어주면 됩니다.
int par[10020];
int fnd(int a){
if (par[a] == a){ return a; }
return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }
void Main(){
int n, m; cin >> n >> m;
for (int i = 1; i <= n; i++){ par[i] = i; }
priority_queue< pi3, vector<pi3>, greater<pi3> > pq;
while (m--){
int v, w, x; cin >> v >> w >> x;
pq.push({ x, {v, w} });
}
int ans = 0;
while (!pq.empty()){
int x = pq.top().fr;
int v = pq.top().sc.fr, w = pq.top().sc.sc;
pq.pop();
if (fnd(v) == fnd(w)){ continue; }
uni(v, w); ans += x;
}
cout << ans;
}MST 연습 문제입니다. Prim으로 풀어주면 됩니다.
bool chk[10020];
vector<pi2> adj[10020];
void Main(){
int n, m; cin >> n >> m;
while (m--){
int v, w, x; cin >> v >> w >> x;
adj[v].push_back({w, x}); adj[w].push_back({v, x});
}
priority_queue< pi2, vector<pi2>, greater<pi2> > pq;
chk[1] = 1; for (pi2 p : adj[1]){ pq.push({p.sc, p.fr}); }
int ans = 0;
while (!pq.empty()){
int d = pq.top().fr, v = pq.top().sc; pq.pop();
if (chk[v]){ continue; } chk[v] = 1; ans += d;
for (pi2 p : adj[v]){ pq.push({p.sc, p.fr}); }
}
cout << ans;
}5. [1647] 도시 분할 계획 (Gold IV)
MST를 구하되, 두 마을로 나눌 수 있으니 \( N-2 \)개의 간선만 사용하면 됩니다.
int par[100020];
int fnd(int a){
if (par[a] == a){ return a; }
return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }
void Main(){
int n, m; cin >> n >> m;
for (int i = 1; i <= n; i++){ par[i] = i; }
priority_queue< pi3, vector<pi3>, greater<pi3> > pq;
while (m--){
int v, w, x; cin >> v >> w >> x;
pq.push({ x, {v, w} });
}
int cnt = n-2, ans = 0;
while (cnt){
int x = pq.top().fr;
int v = pq.top().sc.fr, w = pq.top().sc.sc;
pq.pop();
if (fnd(v) == fnd(w)){ continue; }
uni(v, w); ans += x; cnt -= 1;
}
cout << ans;
}6. [14621] 나만 안되는 연애 (Gold III)
MST를 구하면 되지만, '남초' - '여초'를 연결하는 간선들만 남겨야 합니다.
또한, MST가 정말 모든 정점을 포함하고 있는지도 확인해줘야 합니다.
그러니 그렇게 해주면 됩니다.
이 문제에서 '경로'가 실제 두 정점간의 경로와는 관련없음에 주의하세요.
int par[1020];
int fnd(int a){
if (par[a] == a){ return a; }
return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }
bool chk[1020];
void Main(){
int n, m; cin >> n >> m;
for (int i = 1; i <= n; i++){ char ch; cin >> ch; chk[i] = ch=='M'; }
for (int i = 1; i <= n; i++){ par[i] = i; }
priority_queue< pi3, vector<pi3>, greater<pi3> > pq;
while (m--){
int v, w, x; cin >> v >> w >> x;
if (chk[v] == chk[w]){ continue; }
pq.push({ x, {v, w} }); //cout << v << ' ' << w << " / " << x << endl;
}
int cnt = n-1, ans = 0;
while (!pq.empty() && cnt){
int x = pq.top().fr;
int v = pq.top().sc.fr, w = pq.top().sc.sc;
pq.pop();
if (fnd(v) == fnd(w)){ continue; }
uni(v, w); ans += x; cnt -= 1;
}
cout << (cnt ? -1 : ans);
}7. [2887] 행성 터널 (Gold I)
간선의 개수가 \( O(N^2) \)이 되는, 무서운 문제입니다.
하지만 MST에 포함될 수 있는 후보는 "한 축을 기준으로 정렬한 뒤, 인접한 도시들을 잇는 간선"들밖에 되지 않습니다.
만약 그렇지 않다고 합시다. 즉, 어느 축으로도 정렬해도 인접하지 않지만 MST에 "포함되어야만 하는" 간선 \( v \leftrightarrow w \)가 존재한다고 합시다.
정의에 의해, 어떤 축으로 정렬했을 때 \( v \)와 \( w \) 사이에 \( x \)라는 정점이 존재해야 합니다.
그럼, \( v \leftrightarrow x \leftrightarrow w \)의 두 간선을 사용할 수 있게 되고, \( v \le x \le w \)일 때 |v-x| + |x-w| = |v-w| \)가 되니, \( v \leftrightarrow w \) 대신 \( v \leftrightarrow x \leftrightarrow w \)의 두 간선을 써줄 수 있게 됩니다.
즉, \( v \leftrightarrow w \)가 항상 MST에 들어가야 한다는 가정에 위배됩니다.
한 축을 기준으로 정렬한 뒤, 인접한 도시들을 잇는 간선은 축 하나당 \( O(N) \)개니, 이제 약 \( 3N \)개의 간선으로 MST를 구해주면 됩니다.
#define x fr.fr
#define y sc.fr
#define z sc.sc
#define idx fr.sc
int par[100020];
int fnd(int a){
if (par[a] == a){ return a; }
return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }
pl4 pos[100020];
inline ll dis(pl4 p1, pl4 p2){ return min({ abs(p1.x-p2.x), abs(p1.y-p2.y), abs(p1.z-p2.z) }); }
void Main(){
int n; cin >> n;
for (int i = 1; i <= n; i++){ cin >> pos[i].x >> pos[i].y >> pos[i].z; pos[i].idx = i; }
for (int i = 1; i <= n; i++){ par[i] = i; }
priority_queue< pl3, vector<pl3>, greater<pl3> > pq;
sort(pos+1, pos+n+1, [](pl4 a, pl4 b){ return a.x < b.x; });
for (int i = 1; i < n; i++){ pq.push({ dis(pos[i], pos[i+1]), {pos[i].idx, pos[i+1].idx} }); }
sort(pos+1, pos+n+1, [](pl4 a, pl4 b){ return a.y < b.y; });
for (int i = 1; i < n; i++){ pq.push({ dis(pos[i], pos[i+1]), {pos[i].idx, pos[i+1].idx} }); }
sort(pos+1, pos+n+1, [](pl4 a, pl4 b){ return a.z < b.z; });
for (int i = 1; i < n; i++){ pq.push({ dis(pos[i], pos[i+1]), {pos[i].idx, pos[i+1].idx} }); }
int cnt = n-1, ans = 0;
while (!pq.empty() && cnt){
int d = pq.top().fr;
int v = pq.top().sc.fr, w = pq.top().sc.sc;
pq.pop();
if (fnd(v) == fnd(w)){ continue; }
uni(v, w); ans += d; cnt -= 1;
}
cout << (cnt ? -1 : ans);
}8. [20040] 사이클 게임 (Gold IV)
MST 구하던대로 해주면 됩니다.
간선을 하나씩 입력받으면서 잇다가 이미 두 정점이 같은 집합에 속해있을 때 답을 출력하고 종료해주면 됩니다.
만약 위 경우가 끝까지 나오지 않는다면 0을 출력하고 종료해야 함에 주의하세요.
int par[500020];
int fnd(int a){
if (par[a] == a){ return a; }
return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }
void Main(){
int n, m; cin >> n >> m;
for (int i = 0; i < n; i++){ par[i] = i; }
for (int ans = 1; ans <= m; ans++){
int v, w; cin >> v >> w;
if (fnd(v) == fnd(w)){ cout << ans; return; }
uni(v, w);
}
cout << 0;
}9. [14595] 동방 프로젝트 (Large) (Gold II)
Union 연산을 구간에 한 번에 넣는 문제입니다.
이는 집합의 루트를 맨 오른쪽에 놓는 방식으로 빠르게 해결할 수 있으며, 이 이유는
집합의 루트를 찾으면 그 바로 오른쪽 정점은 항상 다른 집합에 속하기 때문입니다.
int par[1000020];
int fnd(int a){
if (par[a] == a){ return a; }
return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }
void Main(){
int n; cin >> n;
for (int i = 1; i <= n; i++){ par[i] = i; }
int q; cin >> q; while (q--){
int st, ed; cin >> st >> ed;
int ptr = fnd(st); while (ptr < ed){
uni(ptr, ptr+1); ptr = fnd(ptr+1);
}
}
int ans = 0;
for (int i = 1; i <= n; i++){ ans += i==fnd(i); }
cout << ans;
}10. [17132] 두더지가 정보섬에 올라온 이유 (Platinum V)
가중치가 큰 간선부터 이으면서 양쪽 집합에 생기는 새로운 경로의 수 (= 두 집합의 크기의 곱)만큼 답에 더해주면 됩니다.
즉, 간선을 이으면서 새로 생기는 모든 경로는 정확히 그 가중치와 동일한 행복도를 지니게 됩니다.
이유는 간단한데, 지금까지 이은 모든 간선은 이 간선보다 가중치가 크거나 같기 때문입니다.
ll ans = 0;
int par[100020], siz[100020];
int fnd(int a){
if (par[a] == a){ return a; }
return par[a] = fnd(par[a]);
}
void uni(int a, int b, ll x){ int pa = fnd(a), pb = fnd(b);
ans += x*siz[pa]*siz[pb];
par[pa] = pb; siz[pb] += siz[pa];
}
void Main(){
int n; cin >> n;
for (int i = 1; i <= n; i++){ par[i] = i; siz[i] = 1; }
priority_queue<pi3> pq;
for (int i = 1; i < n; i++){
int v, w, x; cin >> v >> w >> x;
pq.push({ x, {v, w} });
}
while (!pq.empty()){
int x = pq.top().fr;
int v = pq.top().sc.fr, w = pq.top().sc.sc;
pq.pop(); uni(v, w, x);
}
cout << ans;
}11. [13306] 트리 (Platinum IV)
쿼리를 주어진 순서대로 처리하기에는 너무 어려운 문제입니다.
하지만 이를 역순으로 처리한다면?
트리의 최종 상태인 모든 정점이 다 나눠져 있는 상태에서 시작해서
0번 쿼리로 두 정점을 잇고, 1번 쿼리로 두 정점이 같은 트리 (= 집합)에 속했는지 확인하는 문제가 됩니다.
익숙하지 않나요? 위 문제는 "1. [1717] 집합의 표현"과 정확히 동일한 문제이기 때문입니다!
정확히는, 두 정점을 이을 때 (나, 나의 부모)를 잘 찾아서 잇는다는 점이 다르긴 합니다.
근데 이건 그냥 각 정점별로 부모를 저장하면 간단하게 처리됩니다.
+ 쿼리를 역순으로 처리하니까 출력도 역순으로 해야 함에 주의하세요.
+ Q는 최대 20만이지만, 입력되는 쿼리의 개수는 최대 40만임에 주의하세요.
int par[200020];
int fnd(int a){
if (par[a] == a){ return a; }
return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }
int ptr[200020];
pi3 qry[400020];
void Main(){
int n, q; cin >> n >> q; q += n-1;
for (int i = 1; i <= n; i++){ par[i] = i; }
for (int i = 2; i <= n; i++){ cin >> ptr[i]; }
for (int i = 1; i <= q; i++){
int typ; cin >> typ;
if (typ == 0){ int v; cin >> v; qry[i] = { 0, {v, ptr[v]} }; }
if (typ == 1){ int v, w; cin >> v >> w; qry[i] = { 1, {v, w} }; }
}
stack<bool> ans;
for (int i = q; i >= 1; i--){
int typ = qry[i].fr;
int v = qry[i].sc.fr, w = qry[i].sc.sc;
if (typ == 0){ uni(v, w); }
if (typ == 1){ ans.push( fnd(v) == fnd(w) ); }
}
while (!ans.empty()){
bool res = ans.top(); ans.pop();
cout << (res ? "YES" : "NO") << endl;
}
}12. [17490] 일감호에 다리 놓기 (Gold III)
공사 중이 아니라 서로 건너갈 수 있는 강의동을 하나의 집합으로 묶으면,
자명히 하나의 강의동 집합에서 와우도까지 가는데 필요한 돌의 개수는 \( \min S_{i} \)가 됩니다.
그러니 이 문제도 결국 Union-Find를 하면서 각 집합별로 필요한 돌의 개수의 최솟값을 같이 관리해주면 됩니다.
\( M \le 1 \)이라서 돌 없이도 이동할 수 있는 경우를 예외처리해야 함에 주의하세요.
int par[1000020], cnt[1000020];
int fnd(int a){
if (par[a] == a){ return a; }
return par[a] = fnd(par[a]);
}
void uni(int a, int b){ int pa = fnd(a), pb = fnd(b);
par[pa] = pb; cnt[pb] = min(cnt[pb], cnt[pa]);
}
bool chk[1000020];
void Main(){
int n, m; ll k; cin >> n >> m >> k;
if (m <= 1){ cout << "YES"; return; }
for (int i = 1; i <= n; i++){ par[i] = i; cin >> cnt[i]; }
for (int i = 1; i <= m; i++){
int v, w; cin >> v >> w;
if ((v+1)%n == w%n){ chk[v] = 1; }
if ((w+1)%n == v%n){ chk[w] = 1; }
}
for (int i = 1; i <= n; i++){
if (!chk[i]){
if (i != n){ uni(i, i+1); } else{ uni(n, 1); }
}
}
ll ans = 0;
for (int i = 1; i <= n; i++){
if (fnd(i) != i){ continue; }
ans += cnt[i];
}
cout << (ans <= k ? "YES" : "NO");
}13. [1774] 우주신과의 교감 (Gold III)
이미 연결된 ( = 필수적으로 연결해야 하는) 간선들이 있는 MST입니다.
풀이는 간단한데, 어차피 연결할 거 그냥 시작하기 전에 연결해주고 시작해주면 됩니다.
간선의 가중치 = 두 정점의 거리 이고, 실수오차는 최대한 피하는 걸 추천하는 사람으로서
최종 답인 ans에 더할 때까지는 가중치를 두 정점의 거리의 제곱으로 저장했습니다.
int par[1020];
int fnd(int a){
if (par[a] == a){ return a; }
return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }
inline ll dis(pl2 p1, pl2 p2){
ll y = p1.fr-p2.fr, x = p1.sc-p2.sc;
return y*y + x*x;
}
pl2 pos[1020];
void Main(){
int n, m; cin >> n >> m;
for (int i = 1; i <= n; i++){ par[i] = i; }
for (int i = 1; i <= n; i++){ cin >> pos[i].fr >> pos[i].sc; }
priority_queue< pl3, vector<pl3>, greater<pl3> > pq;
for (int i = 1; i <= n; i++){
for (int j = i+1; j <= n; j++){
pq.push({ dis(pos[i], pos[j]), {i, j} });
}
}
while (m--){ int v, w; cin >> v >> w; uni(v, w); }
ld ans = 0;
while (!pq.empty()){
ll x = pq.top().fr;
int v = pq.top().sc.fr, w = pq.top().sc.sc;
pq.pop();
if (fnd(v) == fnd(w)){ continue; }
uni(v, w); ans += sqrt((ld)x);
}
cout << ans;
}14. [1045] 도로 (Gold I)
Spanning Tree긴 한데, 우선 순위를 최대화하라고 합니다.
그런데 "앞에 있는 간선의 우선순위가 크면" 뒤의 가중치와 상관없이 더 높은 우선순위를 가지기 때문에, 결국 우선순위가 높은 간선부터 차례대로 보면 됩니다.
즉, 우선순위로 간선 재정렬하고 아무 일 없다는 듯이 MST 돌리면 됩니다.
+ 도로를 정확히 M개 가져야 하므로, 두 정점이 같은 집합에 속할 때에도 총 M-(N-1)회는 연결시켜줘야 합니다.
int par[60];
int fnd(int a){
if (par[a] == a){ return a; }
return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }
int ans[60];
void Main(){
priority_queue< pi2, vector<pi2>, greater<pi2> > pq;
int n, m; cin >> n >> m; m -= n-1;
for (int i = 0; i < n; i++){ par[i] = i; }
for (int i = 0; i < n; i++){
for (int j = 0; j < n; j++){
char ch; cin >> ch;
if (ch == 'Y' && i < j){ pq.push({i, j}); }
}
}
int cnt = n-1;
while (!pq.empty()){
int v = pq.top().fr, w = pq.top().sc;
pq.pop();
if (fnd(v) == fnd(w)){ if (m){ ans[v] += 1; ans[w] += 1; m -= 1; } }
else{ uni(v, w); ans[v] += 1; ans[w] += 1; cnt -= 1; }
}
if (cnt || m){ cout << -1; return; }
for (int i = 0; i < n; i++){ cout << ans[i] << ' '; }
}15. [23034] 조별과제 멈춰! (Platinum V)
MST를 구한 뒤 x와 y를 잇는 경로 중 가중치가 가장 큰 간선을 제거하면 됩니다.
x - y를 미리 연결시키고 MST를 돌린다고 생각하면 편합니다.
이렇게만 해도 되는 이유가, 위 과정 때문에 MST에 생기는 변화는 x - y를 잇는 경로의 변화밖에 없기 때문입니다.
그렇지 않은 간선이 있다면, 애초부터 MST를 만들 때에 그 간선을 먼저 볼 것이기 때문입니다.
int par[1020];
int fnd(int a){
if (par[a] == a){ return a; }
return par[a] = fnd(par[a]);
}
void uni(int a, int b){ par[fnd(a)] = fnd(b); }
vector<pi2> adj[1020];
ll dis[1020][1020];
void dfs(int st, int now, int pre, ll mx){
dis[st][now] = mx;
for (pi2 p : adj[now]){
int nxt = p.fr; ll dst = p.sc;
if (nxt == pre){ continue; }
dfs(st, nxt, now, max(mx, dst));
}
}
void Main(){
priority_queue< pi3, vector<pi3>, greater<pi3> > pq;
int n, m; cin >> n >> m;
for (int i = 1; i <= n; i++){ par[i] = i; }
while (m--){
int v, w, x; cin >> v >> w >> x;
pq.push({ x, {v, w} });
}
ll ans = 0;
while (!pq.empty()){
int x = pq.top().fr;
int v = pq.top().sc.fr, w = pq.top().sc.sc;
pq.pop();
if (fnd(v) == fnd(w)){ continue; }
uni(v, w); ans += x;
adj[v].push_back({w, x}); adj[w].push_back({v, x});
}
for (int i = 1; i <= n; i++){ dfs(i, i, -1, 0); }
int q; cin >> q; while (q--){
int v, w; cin >> v >> w;
cout << ans - dis[v][w] << endl;
}
}