public static int getNeighbor(int row, int col, List<List<Integer>> matrix, int[][] index){
        int[][] DIR = {{1, 0}, {-1, 0}, {0 , -1}, {0, 1}};
        int min = Integer.MAX_VALUE;
        int minRow = 0;
        int minCol = 0;
        for(int[] dir: DIR){
            int nRow = row + dir[0];
            int nCol = col + dir[1];
            if(nRow>=0 && nCol>= 0 && nRow < matrix.size() && nCol<matrix.get(0).size()){
                if(matrix.get(nRow).get(nCol) < min){
                    min = matrix.get(nRow).get(nCol);
                    minRow = nRow;
                    minCol = nCol;
                }
            }
        }

        if(min <= matrix.get(row).get(col)){
            return index[minRow][minCol];
        }
        return -1;
    }

    public static int find(int node, Map<Integer, Integer> parent){

        Integer curr = node;
        while(parent.get(curr) != curr){
            parent.put(curr, parent.get(parent.get(curr))); // path compression
            curr = parent.get(curr);
        }
        return curr;
    }

    public static void union(int u, int v, Map<Integer, Integer> parent, Map<Integer, Integer> size, int[][] index){

        int rootu = find(u, parent);
        int rootv = find(v, parent);
        if(size.get(rootu) <= size.get(rootv)){
            parent.put(rootu, rootv);
            size.put(rootv, size.get(rootv) + size.get(rootu));
            size.put(rootu, 0);
        }
        else { // size.get(rootv) < size.get(rootu)
            parent.put(rootv, rootu);
            size.put(rootu, size.get(rootv) + size.get(rootu));
            size.put(rootv, 0);
        }
    }

    public static List<Integer> find_basins(List<List<Integer>> matrix) {
        //Write your code here
        /**
         * Steps:
         * - implement Kruskal's algorithm
         *  - for each element on the matrix figure out its neigbors's minimum value
         *  - find the neighbors's element and do a union
         */

        Map<Integer, Integer> parent = new HashMap<>();
        Map<Integer, Integer> size = new HashMap<>();
        int[][] index = new int[matrix.size()][matrix.get(0).size()];
        List<Integer> results = new ArrayList<>();

        // initialize the parent and size maps
        for(int row=0; row<matrix.size(); row++) {
            for (int col = 0; col < matrix.get(0).size(); col++) {
                int idx = row * matrix.get(0).size() + col;
                index[row][col] = idx;
                parent.put(idx, idx);
                size.put(idx, 1);
            }
        }

        // start going through each cell and performing union
        for(int row=0; row<matrix.size(); row++){
            for(int col=0; col<matrix.get(0).size(); col++){
                int neighbor = getNeighbor(row, col, matrix, index);
                if(neighbor != -1){ // perform union
                    union(index[row][col], neighbor, parent, size, index);
                }
            }
        }

        // construct the results
        for(int i = 0; i <  matrix.size()*matrix.get(0).size(); i++){
            if(i == parent.get(i)){
                results.add(size.get(i));
            }
        }
        Collections.sort(results);
        return results;
    }
}