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

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

[AtCoder] AtCoder Beginner Contest 097 D - Equals

グラフに関する問題

D - Equals

解き方が分からなかったので、解説を見ると、この問題はグラフに関する問題だということが分かりました。解説では、素集合データ構造を使って解くことができるそうですが、僕は隣接リストを作って解くことにしました。

隣接リスト

隣接リストの作り方を下の本で探すと、

プログラミングコンテスト攻略のためのアルゴリズムとデータ構造

プログラミングコンテスト攻略のためのアルゴリズムとデータ構造

  1. 287 の12.5 連結成分に隣接リストを使う問題がありました。問題は、

グラフの連結成分 | アルゴリズムとデータ構造 | Aizu Online Judge

で見ることができます。今回のD問題と上記の問題は非常に似ている問題なので、この連結成分の問題を解くことができれば、D問題も解けるでしょう。

この本ではC++で記述されているので、Javaで書き直す必要があります。

隣接リストの活用

隣接リストを作ることができたら、探索をする必要があります。この本では、深さ優先探索を使って解いているので、その部分も書き直す必要があります。

コード

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Deque;
import java.util.Map;
import java.util.Scanner;
import java.util.TreeMap;

public class ProblemD2 {
    public static Map<Integer, ArrayList<Integer>> map;
    public static int N;
    public static void main(String[] args) {
        Scanner scan = new Scanner(new BufferedReader(new InputStreamReader(System.in)));
        N = scan.nextInt();
        int M = scan.nextInt();
        Integer []p = new Integer[N];
        int []x = new int[M];
        int []y = new int[M];
        for(int i = 0; i < N; i++) {
            p[i] = scan.nextInt();
        }
        for(int i = 0; i < M; i++) {
            x[i] = scan.nextInt();
            y[i] = scan.nextInt();
        }
        scan.close();

        // 隣接リスト
        map = new TreeMap<Integer, ArrayList<Integer>>();
        for(int i = 0; i < M; i++) {
            if(!map.containsKey(x[i])) {
                ArrayList<Integer> list = new ArrayList<Integer>();
                list.add(y[i]);
                map.put(x[i], list);
            }else {
                map.get(x[i]).add(y[i]);
            }
            if(!map.containsKey(y[i])) {
                ArrayList<Integer> list = new ArrayList<Integer>();
                list.add(x[i]);
                map.put(y[i], list);
            }else {
                map.get(y[i]).add(x[i]);
            }
        }

//     for(int i : map.keySet()) {
//         System.out.print(i + " : ");
//         for(int j = 0; j < map.get(i).size(); j++) {
//             System.out.print(map.get(i).get(j) + " ");
//         }
//         System.out.println();
//     }

        int []color = new int[N + 1];
        Arrays.fill(color, 0);

        assignColor(color);
//     for(int i = 1; i <= N; i++) {
//         System.out.println(color[i]);
//     }

        int t = 0;
        int sum = 0;

        for(int i = 1; i <= N; i++) {
//         for(int j = 0; j < N; j++) {
//             if(i == p[j]) {
//                 t = j + 1;
//                 break;
//             }
//         }
            if(color[p[i - 1]] == color[i]) {
                sum ++;
            }
        }

        System.out.println(sum);


    }
    public static void dfs(int r, int c, int[] color) {
        Deque<Integer> stack = new ArrayDeque<Integer>();
        stack.push(r);
        color[r] = c;
        while(!stack.isEmpty()) {
            int u = stack.getFirst();
            stack.pop();
            if(map.containsKey(u)) {
                for(int i = 0; i < map.get(u).size(); i++) {
                    int v = map.get(u).get(i);
                    if(color[v] == 0) {
                        color[v] = c;
                        stack.push(v);
                    }
                }
            }
        }
    }
    public static void assignColor(int[] color) {
        int id = 1;
        for(int i = 1; i <= N; i++) {
            if(color[i] == 0) {
                dfs(i, id++, color);
            }
        }
    }
}