题解 | 【模板】并查集

import java.io.*;
import java.util.*;

public class Main {
    // 存储节点 i 的父节点编号；当 parent[i] == i 时，i 为该集合的根节点
    private static int[] parent;
    // 仅在 i 为根节点时有效，记录以 i 为根的集合中包含的元素总数
    private static int[] size;

    public static void main(String[] args) throws IOException {
        // 使用高性能输入输出流
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        PrintWriter out = new PrintWriter(new OutputStreamWriter(System.out));

        // 解析输入首行的 n (元素总数) 和 q (操作总数)
        String[] strA = br.readLine().trim().split("\\s+");
        int n = Integer.parseInt(strA[0]);
        int q = Integer.parseInt(strA[1]);

        parent = new int[n + 1];
        size = new int[n + 1];
        for (int i = 1; i <= n; i++) {
            // 初始状态下，每个元素自成一个集合，父节点指向自己
            parent[i] = i;
            // 每个初始集合的大小均为 1
            size[i] = 1;
        }

        while (q-- > 0) {
            String[] strB = br.readLine().trim().split("\\s+");
            int op = Integer.parseInt(strB[0]); // 操作类型编号

            switch (op) {
                case 1: {
                    // 合并操作：将 i 和 j 所在的两个集合合并为一个
                    int i = Integer.parseInt(strB[1]);
                    int j = Integer.parseInt(strB[2]);
                    union(i, j);
                    break;
                }
                case 2: {
                    // 判断 i 和 j 是否属于同一个集合
                    int i = Integer.parseInt(strB[1]);
                    int j = Integer.parseInt(strB[2]);
                    // 若两者的根节点相同，则在同一集合，输出 YES
                    out.println(find(i) == find(j) ? "YES" : "NO");
                    break;
                }
                case 3: {
                    // 获取 i 所在集合的元素总数
                    int i = Integer.parseInt(strB[1]);
                    // 必须先通过 find(i) 找到根节点，因为只有根节点的 size 是准确的
                    out.println(size[find(i)]);
                    break;
                }
            }
        }

        // 刷新缓冲区并释放 IO 资源
        out.flush();
        out.close();
        br.close();
    }

    /**
     * 查找元素 i 所在集合的根节点
     * 在递归查找过程中，将路径上所有节点的父节点直接指向根节点
     */
    private static int find(int i) {
        // 如果自己就是根节点，直接返回
        if (parent[i] == i) {
            return i;
        }
        // 将查找到的根节点赋值给当前节点的 parent 数组
        return parent[i] = find(parent[i]);
    }

    /**
     * 将包含 i 和 j 的两个集合合并
     *
     * @param i 第一个元素
     * @param j 第二个元素
     */
    private static void union(int i, int j) {
        // 分别获取两个元素当前的根节点
        int rootI = find(i);
        int rootJ = find(j);

        // 如果根节点不同，说明处于不同集合，执行合并
        if (rootI != rootJ) {
            // 将 I 集合的根节点挂载到 J 集合的根节点之下
            parent[rootI] = rootJ;
            // 将被合并集合的规模累加到新的根节点 (rootJ) 上
            size[rootJ] += size[rootI];
        }
    }
}