个人练习-PAT甲级-1118 Birds in Forest

题目链接https://pintia.cn/problem-sets/994805342720868352/problems/994805354108403712

题目大意：给出N张图片，每张图片上有K只鸟，同一张图片内的鸟视为在同一棵树上。求有几棵树，几只鸟。并且给出Q个query，每个query给出两只鸟编号，判断它们是否属于同一棵树。

和之前那道【1114 Family Property】很像，也是并查集问题。由于用数组建set的话数元素个数有点麻烦（虽然也可以实现），干脆直接用STL的set了。思路几乎不变，依旧是注意什么时候将两个set合并。我们建议将编号大的set合并到编号小的set里

            for (int t = 0; t < arr.size(); t++) {
    if (arr[t].count((B))) {
    if (root == -1)root = t;else if (root != t) {
    int big = max(root, t), small = min(root, t);Union(small, big);root = small;}break;}}

不过最后处理query时，好像没有数组方便。。。什么时候有空再用数组写一遍吧

        for (int j = 0; j < arr.size(); j++) {
    if (arr[j].count(b1) && !arr[j].count(b2)) {
    printf("No\n");break;}else if (!arr[j].count(b1) && arr[j].count(b2)) {
    printf("No\n");break;}else if (arr[j].count(b1) && arr[j].count(b2)) {
    printf("Yes\n");break;}}

完整代码

#include <iostream>
#include <cstdio>
#include <cmath>
#include <vector>
#include <algorithm>
#include <queue>
#include <map>
#include <set>using namespace std;vector<set<int>> arr;void Union(int small, int big) {
    for (auto it = arr[big].begin(); it != arr[big].end(); it++)arr[small].insert(*it);arr.erase(arr.begin()+big, arr.begin()+big+1);
}int main() {
    int N, K, Q, B;scanf("%d", &N);for (int i = 0; i < N; i++) {
    int root = -1;scanf("%d", &K);vector<int> birds(K);for (int j = 0; j < K; j++) {
    scanf("%d", &B);birds[j] = B;// find B_k's belongingfor (int t = 0; t < arr.size(); t++) {
    if (arr[t].count((B))) {
    if (root == -1)root = t;else if (root != t) {
    int big = max(root, t), small = min(root, t);Union(small, big);root = small;}break;}}}if (root == -1) {
    root = arr.size();set<int> tmp;for (int j = 0; j < K; j++)tmp.insert(birds[j]);arr.push_back(tmp);}else {
    for (int j = 0; j < K; j++)arr[root].insert(birds[j]);}}int cnt = 0;for (int i = 0; i < arr.size(); i++)cnt += arr[i].size();printf("%d %d\n", arr.size(), cnt);scanf("%d", &Q);for (int i = 0; i < Q; i++) {
    int b1, b2;scanf("%d %d", &b1, &b2);for (int j = 0; j < arr.size(); j++) {
    if (arr[j].count(b1) && !arr[j].count(b2)) {
    printf("No\n");break;}else if (!arr[j].count(b1) && arr[j].count(b2)) {
    printf("No\n");break;}else if (arr[j].count(b1) && arr[j].count(b2)) {
    printf("Yes\n");break;}}}return 0;
}