Find the Weak Connected Component in the Directed Graph

public class Solution {
    /**
     * @param nodes a array of Directed graph node
     * @return a connected set of a directed graph
     */
     
    class UnionFind {
        HashMap<Integer, Integer> map = new HashMap<Integer, Integer>();
        public int find(int a) {
            if (!map.containsKey(a)) {
                map.put(a, a);
                return a;
            }
            int copy = a;
            while (copy != map.get(copy)) {
                copy = map.get(copy);
            }
            while (a != map.get(a)) {
                int temp = map.get(a);
                map.put(a, copy);
                a = temp;
            }
            return a;
        }
    }
    public List<List<Integer>> connectedSet2(ArrayList<directedgraphnode> nodes) {
        // Write your code here
        List<List<Integer>> result = new ArrayList<List<Integer>>();
        if (nodes == null || nodes.size() == 0) return result;
        
        UnionFind uf = new UnionFind();
        for (DirectedGraphNode node : nodes) {
            for (DirectedGraphNode neighbor : node.neighbors) {
                int father = uf.find(node.label);
                int son = uf.find(neighbor.label);
                uf.map.put(father, son);
            }
        }
        
        HashMap<Integer, List<Integer>> graph = new HashMap<Integer, List<Integer>>();
        for (DirectedGraphNode node : nodes) {
            List<Integer> list = new ArrayList<Integer>();
            int root = uf.find(node.label);
            if (graph.containsKey(root)) {
                list = graph.get(root);
            }
            list.add(node.label);
            graph.put(root, list);
        }
        
        for (int key : graph.keySet()) {
            result.add(graph.get(key));
        }
        return result;
    }
}