5-4: 트리에서의 노드 간 거리

트리에서의 거리 1

 

문제

---

트리가 주어지고, 두 노드 X, Y가 주어질 때, 이 두 노드 사이의 거리를 출력하는 프로그램을 작성하시오. 트리에서는 두 노드를 잇는 경로가 유일하기 때문에, 정답은 항상 유일하다는 것을 참고한다. 예를 들어, 다음과 같은 트리에서 노드 3, 노드 6 사이의 거리는 4이다.

 

입력

---

첫 번째 줄에 트리의 노드 개수 n, 두 노드 X, Y의 번호가 주어진다. ( 1 ≤ X, Y ≤ n ≤ 1000 ) 두 번째 줄부터 트리의 간선 정보가 주어진다. 각 줄은 2개의 숫자 a, b로 이루어지며, 이는 노드 a가 노드 b의 부모노드라는 것을 의미한다. 루트는 노드 0이라고 가정한다.  

출력

---

두 노드 X, Y 사이의 거리를 출력한다.

 

예제 입력

11 3 6
0 1
0 2
1 3
1 4
1 5
2 6
2 10
6 7
6 8
6 9

예제 출력

4

import java.util.ArrayList;
import java.util.Scanner;

public class TreeNodeDistance {


    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);

        int N = sc.nextInt();
        TreeNode[] nodes = new TreeNode[N];

        int X = sc.nextInt();
        int Y = sc.nextInt();

        for (int i=0; i<N-1; i++){
            int curData = sc.nextInt();
            int childData = sc.nextInt();

            if (nodes[curData] == null) {
                TreeNode node = new TreeNode(curData);
                nodes[curData] = node;
            }

            if (nodes[childData] == null) {
                TreeNode node = new TreeNode(childData);

                nodes[childData] = node;
                nodes[curData].children.add(nodes[childData]);
            }

        }


        System.out.println(getDistance(nodes[0], nodes[X], nodes[Y]));

    }

    private static int getDistance(TreeNode node, TreeNode node1, TreeNode node2) {
//        System.out.println("distance between root and node1: " + getDistanceH(node, node1));
        return getDistanceH(node, node1) + getDistanceH(node, node2) - 2*(getDistanceH(node, lca(node, node1, node2)));
    }

    private static int getDistanceH(TreeNode node, TreeNode targetNode) {
        if (node == null){
            return 0;
        }

        if (node.children == null){
            return 0;
        }

        if (node.children.contains(targetNode)){
            return 1;
        }

        for (TreeNode child : node.children){
            //call on subchildren.
            int childDistance = getDistanceH(child, targetNode);
            if (childDistance != 0) return 1+ childDistance;
        }

        return 0;

    }

    private static TreeNode lca(TreeNode root, TreeNode x, TreeNode y) {
        if (root == null) {
            return null;
        }
        if (root == x || root ==y){
            return root;
        }

        TreeNode[] lcaNodes = new TreeNode[root.children.size()];
        int index = 0;
        int count = 0;
        TreeNode temp = null;
        for (TreeNode child : root.children){
            lcaNodes[index] = lca(child, x, y);
            if (lcaNodes[index] != null){
                count++;
                temp = lcaNodes[index];
            }
            index++;
        }

        if (count >=2) return root;

        return temp;
    }


    public static class TreeNode{
        int data;
        //        TreeNode left;
////        TreeNode right;
        ArrayList<TreeNode> children;

        public TreeNode(int data){
            this.data = data;
            children = new ArrayList<>();
        }
    }
}