Graph Valid Tree(图是否是树)

http://www.lintcode.com/zh-cn/problem/graph-valid-tree/

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

public class Solution {
    /*
     * @param n: An integer
     * @param edges: a list of undirected edges
     * @return: true if it's a valid tree, or false
     */
    public boolean validTree(int n, int[][] edges) {
        // write your code here
        if (n == 1) {
            return true;
        }
        Map> map = new HashMap<>();
//        map是个森林，最后要成树，要求size为1，并且第一个的节点数为n
        for (int i = 0; i < edges.length; i++) {
            int[] edge = edges[i];
//            取得每条边去查找两个节点在森林中的位置
            Integer key1 = getKey(map, edge[0]);
            Integer key2 = getKey(map, edge[1]);
            if (key1 == null && key2 == null) {
//                如果两个都不存在，它们本身形成一个树。
                Set set = new HashSet<>();
                set.add(edge[0]);
                set.add(edge[1]);
                map.put(edge[0], set);
            } else if (key1 == null) {
//                如果有一个存在，把另外一个节点加入这个树
                map.get(key2).add(edge[0]);
            } else if (key2 == null) {
//                如果有一个存在，把另外一个节点加入这个树
                map.get(key1).add(edge[1]);
            } else {
//                如果两个在同一个树上，说明它们形成了环，成为图，返回
                if (key1 == key2) {
                    return false;
                } else {
//                    如果在两个树上，把这两个树合并
                    Set set1 = map.get(key1);
                    Set set2 = map.get(key2);
                    set2.addAll(set1);
                    map.remove(key1);
                }
            }
        }
        if (map.size() != 1) {
            return false;
        }
        return map.get(map.keySet().iterator().next()).size() == n;
    }


    private Integer getKey(Map> map, int i) {
        for (Integer key : map.keySet()) {
            if (map.get(key).contains(i)) {
                return key;
            }
        }
        return null;
    }
}