LeetCode -1631 最小体力消耗路径

class Solution {
    public class Edge{
        int len;
        int x;
        int y;
        public Edge(int len,int x,int y){
            this.len = len;
            this.x = x;
            this.y = y;
        }
    }
    public int minimumEffortPath(int[][] heights) {
        int m = heights.length;
        int n = heights[0].length;
        List<Edge> edges = new ArrayList<>();
        for(int i=0;i<m;i++){
            for(int j=0;j<n;j++){
                if(i < m-1){
                   int num1 = i * n + j;
                   int num2 = (i+1) * n + j;
                   edges.add(new Edge(Math.abs(heights[i][j] - heights[i+1][j]),num1,num2));
                }
                if(j < n-1){
                   int num1 = i * n + j;
                   int num2 = i * n + j + 1;
                   edges.add(new Edge(Math.abs(heights[i][j] - heights[i][j+1]),num1,num2));
                }
            }
        }
        Collections.sort(edges,new Comparator<Edge>(){
            public int compare(Edge edge1,Edge edge2){
                return edge1.len - edge2.len;
            }
        });
        UnionFind uf = new UnionFind(m*n);
        int num1 = 0;
        int num2 = m*n -1;
        int index = 0;
        Edge edge = new Edge(0,0,0);
        while(uf.find(num1)!=uf.find(num2)){
           edge = edges.get(index);
           uf.Union(edge.x,edge.y);
           index += 1;
        }
        return edge.len;

    }
    public class UnionFind{
        int[] parent;
        int count;
        public int getCount(){
            return count;
        }
        public UnionFind(int n){
            parent = new int[n];
            count = n;
            for(int i =0;i<n;i++)
              parent[i] = i;
        }
        public int find(int x){
            if(parent[x] != x)
               parent[x] = find(parent[x]);
            return parent[x];
        }
        public boolean Union(int x,int y){
            int rootX = find(x);
            int rootY = find(y);
            if(rootX == rootY)
                return false ;
            parent[rootX] = rootY;
            count -= 1;
            return true; 
        }
    }
}