[백준/Java] 21606번 : 아침 산책

문제

https://www.acmicpc.net/problem/21606

N개의 노드와 N-1개의 간선으로 이루어진 그래프가 있다.

각 노드는 실내(1) 또는 실외(0)에 해당한다.

출발 노드, 도착 노드가 실내이고 그 사이에 존재하는 노드들이 모두 실외인 경로가 몇 개인지 구하는 문제이다.

풀이

처음 아이디어

처음에는 간선의 양쪽이 실내/실외로 다른 경우 간선 가중치를 1로 두고

간선의 양쪽이 모두 실내인 경우 간선 가중치를 0으로 둬서

모든 경로를 찾은 후 경로의 가중치 합이 0 또는 2가 되면 되지 않나...

라는 아이디어를 떠올렸다. 그리고 실내를 시작점으로 두고 bfs 탐색 하는 방식으로 찾으면 될 것 같았다.

 

그러나 N의 최대값이 너무 크기 때문에 (2<=N<=2*10^5) 

모든 경로를 3초 안에 브루트포스로 탐색할 수 없다.

코드

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

public class Main {
    static int n;
    static List<List<Integer>> graph;
    static int[] location;
    static boolean[] visited;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;

        n = Integer.parseInt(br.readLine());
        String str = br.readLine();

        location = new int[n+1];
        for(int i=0; i<n; i++){
            location[i+1] = str.charAt(i)-48;
        }

        graph = new ArrayList<>();
        for(int i=0; i<n+1; i++){
            graph.add(new ArrayList<>());
        }

        for(int i=0; i<n-1; i++){
            st = new StringTokenizer(br.readLine());
            int x = Integer.parseInt(st.nextToken());
            int y = Integer.parseInt(st.nextToken());
            graph.get(x).add(y);
            graph.get(y).add(x);
        }

        int result = 0;

        visited = new boolean[n+1];
        Arrays.fill(visited, false);
        for(int i=1; i<n+1; i++){
            if(location[i]==0 && !visited[i]){ //실외 찾기
                int k = bfs(i, 0); //실내 개수 리턴
                result += k*(k-1);
                //System.out.println("실외 bfs 결과 실내 개수:"+k);
            }
        }

        visited = new boolean[n+1];
        Arrays.fill(visited, false);
        for(int i=1; i<n+1; i++){
            if(location[i]==1 && !visited[i]){ //실내 찾기
                int k = bfs(i, 1); //실내 개수 리턴
                result += (k-1)*2; //(실내 노드 수 - 1)*2
                //System.out.println("실내 bfs 결과 실내 개수:"+k);
            }
        }

        System.out.println(result);

    }
    static int bfs(int start, int target){
        Queue<Integer> q = new LinkedList<>();

        int cnt = 0;
        int inside = 0;

        q.offer(start);
        visited[start] = true;

        while(!q.isEmpty()){
            int now = q.poll();
            for(int i=0; i<graph.get(now).size(); i++) {
                int next = graph.get(now).get(i);
                if (!visited[next]) {
                    if (location[next] == target) {
                        //아직 방문하지 않았고 찾는 값이 경우에 따라 실외 or 실내이면 큐에 넣기
                        visited[next] = true;
                        q.offer(next);
                        if(target==1 && location[next]==1)
                            inside++;
                    } else { //실내 or 실외 개수 카운트
                        cnt++;
                    }
                }
            }
        }
        if(target==0){
            return cnt;
        }else{
            return inside+1;
        }
    }
}