题解 | #树上最短链#

import java.util.*;
public class Main{
public static void main(String[] args){
Scanner scanner = new Scanner(System.in);
int n = scanner.nextInt();
int[] priorities = new int[n];
for(int i =0;i<n;i++){
priorities[i] = scanner.nextInt();
}
int[][] graph = new int[n-1][2];
for(int i =0;i<graph.length;i++){
for(int j =0;j<2;j++){
graph[i][j] = scanner.nextInt();
}
}
int res = getRes(graph,priorities);
System.out.println(res);
}
public static int getRes(int[][] graph,int[] priorities){
HashMap<Integer,HashSet<node>> p_map = new HashMap<>();
HashMap<Node,List<node>> maps = new HashMap<>();
ArrayList<node> allNodes = new ArrayList<>();
for(int i = 1;i<=graph.length+1;i++){
allNodes.add(new Node(i));
}
int p_a;
int p_b;
Node a;
Node b;
for(int i = 0;i<graph.length;i++){
a = allNodes.get(graph[i][0]-1);
b = allNodes.get(graph[i][1]-1);
p_a = priorities[graph[i][0]-1];
p_b = priorities[graph[i][1]-1];
p_map.computeIfAbsent(p_a,f->new HashSet<>()).add(a);
p_map.computeIfAbsent(p_b,f->new HashSet<>()).add(b);
maps.computeIfAbsent(a,f->new ArrayList<>()).add(b);
maps.computeIfAbsent(b,f->new ArrayList<>()).add(a);
}
Stack<node> stack = new Stack<>();
int ans = Integer.MAX_VALUE;
int res,cur_step;
Node cur;
res = Integer.MAX_VALUE;
for(HashSet<node> set:p_map.values()){
if(set.size()<2){continue;}
ans = Integer.MAX_VALUE;
for(Node node:set){
node.steps = 0;
stack.add(node);
}
while(!stack.isEmpty()){
cur = stack.pop();
cur_step = cur.steps;
for(Node next:maps.get(cur)){
if(next == cur.before){continue;}
if(next.steps != -1){
ans = cur_step + 1 + next.steps;
break;
}
next.before = cur;
next.steps = cur_step+1;
stack.push(next);
}
cur.steps = -1;
cur.before = null;
if(ans != Integer.MAX_VALUE){
break;
}
}
while(!stack.isEmpty()){
cur = stack.pop();
cur.before = null;
cur.steps = -1;
}
if(ans != Integer.MAX_VALUE){
res = Math.min(res,ans);
}
for(Node node:set){
node.steps = -1;
}
}
return res == Integer.MAX_VALUE?-1:res;
}
static class Node{
int steps = -1;
int index;
Node before = null;
public Node(int index){
this.index = index;
}
}
}</node></node></node></node></node>