题解 | #树上最短链#

java版——用深度优先遍历和广度优先遍历求最短路径

深度优先遍历

import java.util.*;

public class Main {
    public static void main(String[] args) {
        Scanner cin = new Scanner(System.in);
        int n = cin.nextInt();
        int[] grade = new int[n + 1];
        for (int i = 1; i <= n; i++) {
            grade[i] = cin.nextInt();
        }
        List<List<Integer>> edges = new ArrayList<>();
        for (int i = 0; i <= n; i++) {
            edges.add(new ArrayList<Integer>());
        }
        for (int i = 0; i < n - 1; i++) {
            int u = cin.nextInt();
            int v = cin.nextInt();
            // 双向边
            edges.get(u).add(v);
            edges.get(v).add(u);
        }
        int minDist = Integer.MAX_VALUE;
        for (int i = 1; i <= n; i++) {
            int currDist = dfs(edges, i, grade);
            if (currDist != -1) {
                minDist = Math.min(minDist, currDist);
            }
        }
        if (minDist == Integer.MAX_VALUE) {
            System.out.println(-1);
        } else {
            System.out.println(minDist);
        }
    }

	// 深度优先
    private static int dfs(List<List<Integer>> edges, int start, int[] grade) {
        int size = grade.length;
        boolean[] visited = new boolean[size];
        Deque<Integer> stack = new LinkedList<>();
        stack.push(start);
        visited[start] = true;
        int dist = 0;
        while (!stack.isEmpty()){
            // 这里先不要出栈，因为待会还要回来遍历它的其它孩子结点
            int u = stack.peek();
            if (u != start && grade[u] == grade[start]){
                return dist;
            }
            for(Integer v : edges.get(u)){
                if (!visited[v]){
                    stack.push(v);
                    visited[v] = true;
                    dist++;
                    break;
                }
            }
            // 栈顶结点和当前结点相同，说明上面循环没有找到它的孩子结点
            // 即该结点的所有孩子结点都已遍历过，则出栈
            if (stack.peek() == u){
                stack.pop();
                dist--;
            }
        }
        return -1;
    }
}

广度优先遍历

思想：用变量 currLevelNodeNum 存储当前层所有结点的个数，然后每出队一个结点， currLevelNodeNum 就减 1 ，当currLevelNodeNum == 0 时说明当前层的所有结点都已出队，则路径长度加一，并用 currLevelNodeNum 记录下一层的结点数，而下一层的结点数就为当前队列中的元素个数。

import java.util.*;

public class Main {
    public static void main(String[] args) {
        Scanner cin = new Scanner(System.in);
        int n = cin.nextInt();
        int[] grade = new int[n + 1];
        for (int i = 1; i <= n; i++) {
            grade[i] = cin.nextInt();
        }
        List<List<Integer>> edges = new ArrayList<>();
        for (int i = 0; i <= n; i++) {
            edges.add(new ArrayList<Integer>());
        }
        for (int i = 0; i < n - 1; i++) {
            int u = cin.nextInt();
            int v = cin.nextInt();
            // 双向边
            edges.get(u).add(v);
            edges.get(v).add(u);
        }
        int minDist = Integer.MAX_VALUE;
        for (int i = 1; i <= n; i++) {
            int currDist = bfs(edges, i, grade);
            if (currDist != -1) {
                minDist = Math.min(minDist, currDist);
            }
        }
        if (minDist == Integer.MAX_VALUE) {
            System.out.println(-1);
        } else {
            System.out.println(minDist);
        }
    }

	// 广度优先
    private static int bfs(List<List<Integer>> edges, int start, int[] grade) {
        int size = grade.length;
        boolean[] visited = new boolean[size];
        Deque<Integer> queue = new LinkedList<>();
        queue.offer(start);
        visited[start] = true;
        int currLevelNodeNum = 1;
        // 路径长度
        int dist = 0;
        while (!queue.isEmpty()) {
            int u = queue.poll();
            currLevelNodeNum--;
            if (u != start && grade[u] == grade[start]) {
                return dist;
            }
            for (Integer v : edges.get(u)) {
                if (!visited[v]) {
                    queue.offer(v);
                    visited[v] = true;
                }
            }
            // 当前层所有结点都已出队
            if (currLevelNodeNum == 0) {
                currLevelNodeNum = queue.size();
                dist++;
            }
        }
        return -1;
    }
}