Graph Valid Tree

Question

from lintcode
Given n nodes labeled from 0 to n - 1 and a list of undirected edges (each edge is a pair of nodes), write a function to check whether these edges make up a valid tree.
Notice
You can assume that no duplicate edges will appear in edges. Since all edges are undirected, [0, 1] is the same as [1, 0] and thus will not appear together in edges.
Example
Given n = 5 and edges = [[0, 1], [0, 2], [0, 3], [1, 4]], return true.

Given n = 5 and edges = [[0, 1], [1, 2], [2, 3], [1, 3], [1, 4]], return false.

Idea

Let me explain the data structure Union Find. This structure is very useful when you only want to find out interconnected domains in a graph. A domain means elements inside are interconnected. For example, in a graph [ a to b, b to c, c to d, e to f], you got 2 domains: a, b, c, d and e, f.

Union Find means you maintain a root for each element. Elements in the same domain must point to the same root element. When two elements connected, you can make one of the two roots points to the other one.

That's all background knowledge needed.

For this question you have to know the definition of a valid tree:

A tree data structure can be defined recursively (locally) as a collection of nodes (starting at a root node), where each node is a data structure consisting of a value, together with a list of references to nodes (the "children"), with the constraints that no reference is duplicated, and none points to the root. -- WikiPedia

The key point is no reference is duplicated.

Also, it is known that a tree with n nodes contains exactly n - 1 edges. This guarantees each node must be linked or link to some nodes.

Think of this: A tree is basically ONE domain. If you keep connecting the input node pairs, they will always point to the same root. BUT if you detect that the pair of nodes already points to the same root before you connect them, you prove the existence of duplicated reference and therefore prove its invalidness.

class UnionFind<T> {

    private HashMap<T,T> rootMap = null;
    private long connectedGroupCount;

    private void addWithoutCheck(T x){
        rootMap.put(x, x);
        connectedGroupCount += 1;
    }

    public UnionFind(){
        rootMap = new HashMap<T, T>();
        connectedGroupCount = 0;
    }

    public void add(T x){
        T root = rootMap.get(x);

        if (root == null){
            addWithoutCheck(x);
        }
    }

    public T find(T x){

        if (!rootMap.containsKey(x)){
            addWithoutCheck(x);
        }
        
        T root = rootMap.get(x);

        if (root.equals(x)){
            return x;
        }
        else{
            rootMap.put(x,find(root));
        }
        return rootMap.get(x);
    }

    public void connect(T a, T b){
        T root_a = find(a);
        T root_b = find(b);

        if (!root_a.equals(root_b)){
            rootMap.put(root_a, root_b);
            connectedGroupCount--;
        }
    }

    public long count(){
        return connectedGroupCount;
    }

}

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) {
        
        if (n-1 != edges.length)
          return false;
        
        UnionFind<Integer> uf = new UnionFind<Integer>();
          
        for (int edge_j = 0; edge_j < edges.length; edge_j++){
            int a = edges[edge_j][0];
            int b = edges[edge_j][1];
            if (uf.find(a) == uf.find(b))
              return false;
            uf.connect(a, b);
        }
        return true;
    }
}