백준 15681 트리와 쿼리 java

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

그래프가 주어지고. 특정 노드를 루트로 삼아서 트리처럼 관리해야한다.

시간제한이 1초인데 10만개의 정점과 쿼리가 주어질 수 있기에 최악의 상황(트리가 1자로 쭉 이어짐)

에서 매번 트리를 탐색하며 갯수를 세면 시간초과가 될 수 밖에 없다.

 

처음구현했던 단순한 방법1과 좀더 가다듬은 방법2 를 기술한다.

 

방법1.

인접리스트를 저장한 그래프를 만들고, 입력된 루트를 기반으로 그래프를 트리로 재구성한 뒤

루트에서부터 트리를 재귀호출로 탐색하며 서브트리의 크기를 갱신한다.

import java.lang.reflect.Array;
import java.util.*;
import java.io.*;
import java.lang.*;

public class Main {
    /*
    만약에 최대 비효율적인 트리일 때 매번 검사하면 10^5^2가 돼서 100억번이나 연산하게됨.
    쿼리 한번 할때마다 내려가면서 세는데, 그때 각 노드의 자손들의 수를 세서 업데이트.

     */
    static class treenode{
        public ArrayList<Integer> children;
        public int subsize;
        public int parent;
        public int num;
        treenode(int num,int parent){
            children=new ArrayList<>();
            subsize=0;
            this.parent=parent;
            this.num=num;
        }
        treenode(){};
    }
    static ArrayList<ArrayList<Integer>> graph = new ArrayList<>();
    static ArrayList<treenode> tree = new ArrayList<>();
    static int nodecount,root,querycount;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        String str = br.readLine();
        StringTokenizer token = new StringTokenizer(str);
        nodecount=Integer.parseInt(token.nextToken());
        root=Integer.parseInt(token.nextToken());
        querycount=Integer.parseInt(token.nextToken());
        for(int i=0;i<=nodecount;i++){
            graph.add(new ArrayList<>());
            tree.add(new treenode());
        }
        for(int i=0;i<nodecount-1;i++){
            str=br.readLine();
            token = new StringTokenizer(str);
            int start = Integer.parseInt(token.nextToken());
            int end = Integer.parseInt(token.nextToken());
            graph.get(start).add(end);
            graph.get(end).add(start);
        }
        /*
        이제 트리를 만들어야한다. 그래프에서 루트로 지정된 것을 꺼내서
        근접한것들을 모두 children에 넣고. 인스턴스화해서 q에도 넣는다.(레벨순회하듯이)
        q에서 poll하고, 근접리스트에서 parent가 아닌 것들은 모두 children에 넣는다.
         */
        Queue<treenode> treeq = new ArrayDeque<>();
        treeq.add(new treenode(root,-1));
        while(!treeq.isEmpty()){
            int len = treeq.size();
            for(int i=0;i<len;i++){
                treenode temp = treeq.poll();
                tree.set(temp.num,temp);
                ArrayList<Integer> graphnode = graph.get(temp.num);
                for(int j=0;j<graphnode.size();j++){
                    if(graphnode.get(j)!=temp.parent){
                        treeq.add(new treenode(graphnode.get(j),temp.num));
                        temp.children.add(graphnode.get(j));
                    }
                }
            }
        }
        childrencountFunc(root);
        for(int i=0;i<querycount;i++){
            int query=Integer.parseInt(br.readLine());
            System.out.println(tree.get(query).subsize);
        }
    }
    static int childrencountFunc(int nodenum){
        treenode tempTreenode = tree.get(nodenum);
        int count=1;
        for(int i=0;i<tempTreenode.children.size();i++){
            count=count+childrencountFunc(tempTreenode.children.get(i));
        }
        tempTreenode.subsize=count;
        return count;
    }
}

 

방법2

굳이 따로 트리를 재구성할 필요 없다. 인접리스트를 타고가면서 순회하는데,

1.dfs개념으로 visited 관리를 해서 이미 방문한 경우 부모로 판단해 서브트리의 크기를 센다.

2.후위순회하며 인자로 parent를 같이 전달하여 부모만 빼고 인접리스트로 서브트리의 크기를 센다.

 

2번방법의 코드.

import java.lang.reflect.Array;
import java.util.*;
import java.io.*;
import java.lang.*;

public class Main {
    /*

     */
    static ArrayList<Integer>[] list = null;
    static int[] dp =null;
    static int nodeCount,rootNum,queryCount;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        String str = br.readLine();
        StringTokenizer token = new StringTokenizer(str);
        nodeCount= Integer.parseInt(token.nextToken());
        rootNum= Integer.parseInt(token.nextToken());
        queryCount= Integer.parseInt(token.nextToken());
        list=new ArrayList[nodeCount+1];
        dp = new int[nodeCount+1];
        for(int i=0;i<=nodeCount;i++){
            list[i]=new ArrayList<>();
        }
        for(int i=0;i<nodeCount-1;i++){
            str=br.readLine();
            token=new StringTokenizer(str);
            int start=Integer.parseInt(token.nextToken());
            int end=Integer.parseInt(token.nextToken());
            list[start].add(end);
            list[end].add(start);
        }
        traversal(rootNum,-1);
        StringBuilder sb = new StringBuilder();
        for(int i=0;i< queryCount;i++){
            int query = Integer.parseInt(br.readLine());
            sb.append(dp[query]+"\n");
        }
        System.out.println(sb);
    }
    static int traversal(int node,int parent){
        int count=1;
        for(int children : list[node]){
            if(children!=parent){
                count+=traversal(children,node);
            }
        }
        dp[node]=count;
        return count;
    }
}