[백준] 15809 전국시대

15809번: 전국시대

 

boj.kr

풀이

이 문제는 유니온 파인드(서로소 집합)를 사용해서 해결할 수 있는 문제이다. 두 나라가 동맹과 전쟁을 수행할 때마다 나라가 합쳐지는데, 이 때 병력을 처리하는 부분만 신경써주면 어렵지 않게 해결할 수 있다. 전쟁을 했는데 두 나라의 병력이 같은 경우에는 부모를 0으로 해줌으로서 정답을 출력할 때 제외할 수 있도록 처리했다.

코드

C++

#include <iostream>
#include <vector>
#include <set>
#include <algorithm>
using namespace std;

int n, m;
vector<int> root, power;

int find(int x) {
    if (root[x] != x)
        root[x] = find(root[x]);
    return root[x];
}

void merge(int a, int b) {
    a = find(a);
    b = find(b);
    if (a < b) {
        root[b] = a;
        power[a] += power[b];
    }
    else {
        root[a] = b;
        power[b] += power[a];
    }
}

void fight(int a, int b) {
    a = find(a);
    b = find(b);

    if (power[a] > power[b]) {
        power[a] -= power[b];
        root[b] = a;
    }
    else if (power[a] < power[b]) {
        power[b] -= power[a];
        root[a] = b;
    }
    else { // 두 나라의 병력이 같은 경우
        root[a] = 0;
        root[b] = 0;
    }
}

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);

    cin >> n >> m;
    root = vector<int>(n + 1);
    power = vector<int>(n + 1);

    for (int i = 1; i <= n; i++)
        cin >> power[i];

    for (int i = 1; i <= n; i++)
        root[i] = i;

    for (int i = 0; i < m; i++) {
        int o, p, q;
        cin >> o >> p >> q;
        if (o == 1) merge(p, q);
        else fight(p, q);
    }

    set<int> s; // 다른 국가에 합쳐진 경우의 중복을 거르기 위해서 set 사용
    vector<int> ans;
    for (int i = 1; i <= n; i++) {
        int p = find(i);
        if (p != 0) // 부모가 0인 나라는 멸망한 나라이므로 제외
            s.insert(p);
    }
    for (int i : s)
        ans.push_back(power[i]);

    sort(ans.begin(), ans.end());
    cout << s.size() << '\n';
    for (int i : ans)
        cout << i << ' ';
    cout << '\n';
    return 0;
}

Python3

import sys


def find(x):
    if root[x] != x:
        root[x] = find(root[x])
    return root[x]


def union(a, b):
    a = find(a)
    b = find(b)

    if a < b:
        root[b] = a
        power[a] += power[b]
    else:
        root[a] = b
        power[b] += power[a]


def fight(a, b):
    a = find(a)
    b = find(b)
    if power[a] > power[b]:
        power[a] -= power[b]
        root[b] = a
    elif power[a] < power[b]:
        power[b] -= power[a]
        root[a] = b
    else:  # 두 나라의 병력이 같은 경우
        root[a] = 0
        root[b] = 0


n, m = map(int, sys.stdin.readline().rstrip().split())
root = [0 for _ in range(n + 1)]
power = [0 for _ in range(n + 1)]
for i in range(1, n + 1):
    root[i] = i

for i in range(1, n + 1):
    power[i] = int(sys.stdin.readline())
for _ in range(m):
    o, p, q = map(int, sys.stdin.readline().rstrip().split())
    if o == 1:
        union(p, q)
    else:
        fight(p, q)

s = set()  # 다른 국가에 합쳐진 경우의 중복을 거르기 위해서 set 사용
ans = []
for i in range(1, n + 1):
    p = find(i)
    if p != 0:  # 부모가 0인 나라는 멸망한 나라이므로 제외
        s.add(find(i))

for i in s:
    ans.append(power[i])
ans.sort()
print(len(s))
for i in ans:
    print(i, end=' ')