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

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

[yukicoder] No.92 逃走経路

問題

No.92 逃走経路 - yukicoder

解説を見ても良く分からなかったので、提出されているコードを見て理解できました。

動的計画法が苦手なのかもしれないです。

考え方

二次元配列 dp_{i, j} を用意します。この配列 dpは、街 iにおいて j回目に犯人が料金所を通過したことを表しています。

料金所の値 d の情報から、 d_1 = c_iとなる i に対して、 a_i または b_i が1回目に犯人がいた可能性のある町となります。

2回目以降は犯人が i - 1回目に存在した可能性のある町で d_i = c_iとなるものを見つけます。

まあ、コードを見た方が分かりやすいですね。

コード

import java.util.Scanner;
import java.util.Set;
import java.util.TreeSet;

public class Exec0092 {
    public static void main(String[] args) {
        Scanner scan = new Scanner(System.in);
        int N = scan.nextInt();
        int M = scan.nextInt();
        int K = scan.nextInt();
        int []a = new int[M];
        int []b = new int[M];
        int []c = new int[M];
        int []d = new int[K];
        for(int i = 0; i < M; i++) {
            a[i] = scan.nextInt();
            b[i] = scan.nextInt();
            c[i] = scan.nextInt();
        }
        for(int i = 0; i < K; i++) {
            d[i] = scan.nextInt();
        }
        scan.close();

        int [][]dp = new int[N + 1][K + 1];
        for(int i = 0; i <= N; i++) {
            for(int j = 0; j <= K; j++) {
                dp[i][j] = 0;
            }
        }

        // 1回目に犯人いる可能性のある街をdp[街][1回目] = 1とする。
        for(int i = 0; i < M; i++) {
            if(c[i] == d[0]) {
                dp[a[i]][1] = 1;
                dp[b[i]][1] = 1;
            }
        }

        // 2回目以降を調べる
        // c[i] = d[i] のとき、i - 1回目に犯人がいた可能性があるとき、i回目に
        // 犯人がいた可能性が存在する。
        for(int i = 1; i < K; i++) {
            for(int j = 0; j < M; j++) {
                if(c[j] == d[i]) {
                    if(dp[a[j]][i] == 1) {
                        dp[b[j]][i + 1] = 1;
                    }
                    if(dp[b[j]][i] == 1) {
                        dp[a[j]][i + 1] = 1;
                    }
                }
            }
        }

        // dp[][K] = 1となっている街を調べる
        Set<Integer> set = new TreeSet<Integer>();
        for(int i = 1; i <= N; i++) {
            if(dp[i][K] == 1) {
                set.add(i);
            }
        }
        System.out.println(set.size());
        int cnt = 0;
        for(int i : set) {
            if(cnt == set.size() - 1) {
                System.out.println(i);
            }else {
                System.out.print(i + " ");
            }
        }
    }
}