부대복귀

사용 알고리즘: bfs, 메모이제이션?

사용 언어: java

 

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

class Solution {
    static int[] dp; //시작점이 destination, 시작점에서 얼마나 떨어져 있는지 계산
    static ArrayList<int[]>[] path; //{갈 수 있는 node번호, 걸리는 시간}
    
    public int[] solution(int n, int[][] roads, int[] sources, int destination) {
        
        //메인
        dp = new int[n+1];
        Arrays.fill(dp, Integer.MAX_VALUE);
        dp[destination] = 0; //자기 자신은 0으로 초기화
        
        path = new ArrayList[n+1];
        for(int i=0; i<n+1; i++) {
            path[i] = new ArrayList<>();
        }
        for(int i=0; i<roads.length; i++) {
            int[] cur = roads[i];
            path[cur[0]].add(new int[] {cur[1], 1});
            path[cur[1]].add(new int[] {cur[0], 1});
        }
        // for(ArrayList<int[]> al : path) {
        //     for(int[] arr : al) System.out.print(Arrays.toString(arr)+" ");
        //     System.out.println();
        // }
        
        boolean[] visit = new boolean[n+1];
        visit[destination] = true;
        bfs(destination, visit);
        // System.out.println(Arrays.toString(dp));
        
        
        int[] answer = new int[sources.length];
        for(int i=0; i<answer.length; i++ ) {
            if(dp[sources[i]] == Integer.MAX_VALUE) answer[i] = -1;
            else answer[i] = dp[sources[i]];
        }
        return answer;
        //메인--
    }
    
    
    private static void bfs(int destination, boolean[] visit) {
        Queue<int[]> q = new ArrayDeque<>();
        q.add(new int[] {destination, 0}); //{현재 위치, 시작점부터 거리}
        
        while(!q.isEmpty()) {
            int[] cur = q.poll();
            
            for(int[] nxt : path[cur[0]]) {
                if( !visit[nxt[0]] ) {
                    visit[nxt[0]] = true;
                    dp[nxt[0]] = cur[1] + 1;
                    q.add(new int[] {nxt[0], cur[1] + 1});
                }
            }
        }
        
    }
    
}

 

처음에 dp를 2차원 배열로 만들려고 했는데, 생각해보니 destination까지의 거리만 구하면 되니까!

bfs도 destination에서 시작하는 경로만 돌면 되고, dp도 1차원 배열로 만들어서 임의의 road부터 destination까지의 거리만 구하면 됐다