山傘のプログラミング勉強日記

プログラミングに関する日記とどうでもよい雑記からなるブログです。

[Aizu Online Judge] Disjoint Set: Union Find Tree (DSL_1_A)

問題

Union-Find木に関する問題です。

互いに素な集合 Union Find| データ構造ライブラリ | Aizu Online Judge

コード

import java.util.Scanner;

public class ProblemA {
    public static void main(String[] args) {
        Scanner scan = new Scanner(System.in);
        int n = scan.nextInt();
        int q = scan.nextInt();
        DisjointSet ds = new DisjointSet(n);

        for(int i = 0; i < q; i++) {
            int com = scan.nextInt();
            int x = scan.nextInt();
            int y = scan.nextInt();
            if(com == 0) {
                ds.unite(x, y);
            }else if(com == 1) {
                if(ds.same(x, y)) {
                    System.out.println(1);
                }else {
                    System.out.println(0);
                }
            }
        }
        scan.close();
    }
    static class DisjointSet {
        int[]p;
        int[]rank;
        DisjointSet(){}
        DisjointSet(int size){
            rank = new int[size];
            p = new int[size];
            for(int i = 0; i < size; i++) {
                makeSet(i);
            }
        }
        void makeSet(int x) {
            p[x] = x;
            rank[x] = 0;
        }
        boolean same(int x, int y) {
            return findSet(x) == findSet(y);
        }
        void unite(int x, int y) {
            link(findSet(x), findSet(y));
        }
        void link(int x, int y) {
            if(rank[x]  > rank[y]) {
                p[y] = x;
            }else {
                p[x] = y;
                if(rank[x] == rank[y]) {
                    rank[y]++;
                }
            }
        }
        int findSet(int x) {
            if(x != p[x]) {
                p[x] = findSet(p[x]);
            }
            return p[x];
        }
    }
}

感想

良く使うデータ構造なのでしっかり覚えたいですね。