[백준]#6059 Pasture Walking

문제

The N cows (2 <= N <= 1,000) conveniently numbered 1..N are grazing among the N pastures also conveniently numbered 1..N. Most conveniently of all, cow i is grazing in pasture i.

Some pairs of pastures are connected by one of N-1 bidirectional walkways that the cows can traverse. Walkway i connects pastures A_i and B_i (1 <= A_i <= N; 1 <= B_i <= N) and has a length of L_i (1 <= L_i <= 10,000).

The walkways are set up in such a way that between any two distinct pastures, there is exactly one path of walkways that travels between them. Thus, the walkways form a tree.

The cows are very social and wish to visit each other often. Ever in a hurry, they want you to help them schedule their visits by computing the lengths of the paths between Q (1 <= Q <= 1,000) pairs of pastures (each pair given as a query p1,p2 (1 <= p1 <= N; 1 <= p2 <= N).

입력

• Line 1: Two space-separated integers: N and Q
• Lines 2..N: Line i+1 contains three space-separated integers: A_i, B_i, and L_i
• Lines N+1..N+Q: Each line contains two space-separated integers representing two distinct pastures between which the cows wish to travel: p1 and p2

출력

• Lines 1..Q: Line i contains the length of the path between the two pastures in query i.

예제 입력 1

4 2
2 1 2
4 3 2
1 4 3
1 2
3 2

예제 출력 1

2
7

풀이

이 문제는 가장 가까운 공통 조상 찾기 알고리즘을 이용해서 구할 수 있었다. 먼저 BFS(너비 우선 탐색) 알고리즘을 이용해서 루트가 1인 트리를 그린 후에 LCA 알고리즘으로 구했다.

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.Queue;

public class Main {
    static int[] parent;
    static int[] depth;
    static ArrayList<Node>[] map;

    public static void main(String[] args) throws Exception{
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        String[] input = br.readLine().split(" ");
        int N = Integer.parseInt(input[0]);
        int Q = Integer.parseInt(input[1]);
        parent = new int[N+1];
        depth = new int[N+1];
        map = new ArrayList[N+1];
        for(int i=1; i<=N; i++)
            map[i] = new ArrayList<>();

        for(int i=0; i<N-1; i++) {
            input = br.readLine().split(" ");
            int start = Integer.parseInt(input[0]);
            int end = Integer.parseInt(input[1]);
            int cost = Integer.parseInt(input[2]);
            map[start].add(new Node(end, cost));
            map[end].add(new Node(start, cost));
        }

        bfs();

        for(int i=0; i<Q; i++) {
            input = br.readLine().split(" ");
            int ans = lca(Integer.parseInt(input[0]), Integer.parseInt(input[1]));

            System.out.println(ans);
        }
    }

    public static int lca(int a, int b) {
        int sum = 0;

        while(true) {
            if(a==b)
                return sum;

            if(depth[a]==depth[b]) {
                while(a!=b) {
                    for(Node n : map[a]) {
                        if(n.child==parent[a])
                            sum += n.cost;
                    }

                    for(Node n : map[b]) {
                        if(n.child==parent[b])
                            sum += n.cost;
                    }

                    a = parent[a];
                    b = parent[b];
                }
            }

            else if(depth[a]<depth[b]) {
                while(depth[a]<depth[b]) {
                    for(Node n : map[b]) {
                        if(n.child==parent[b])
                            sum += n.cost;
                    }

                    b = parent[b];
                }
            }

            else {
                while(depth[a]>depth[b]) {
                    for(Node n : map[a]) {
                        if(n.child==parent[a])
                            sum += n.cost;
                    }

                    a = parent[a];
                }
            }
        }
    }

    public static void bfs() {
        Queue<Node> queue = new LinkedList<>();
        queue.add(new Node(1, 0));
        parent[1] = 1;

        while(!queue.isEmpty()) {
            Node temp = queue.poll();

            for(Node child : map[temp.child]) {
                if(parent[child.child] != 0) continue;

                parent[child.child] = temp.child;
                depth[child.child] = depth[temp.child]+1;
                queue.add(child);
            }
        }
    }

    public static class Node {
        int child;
        int cost;

        public Node(int child, int cost) {
            this.child = child;
            this.cost = cost;
        }
    }
}