Lintcode 178. Graph Valid Tree (Medium) (Python)

Graph Valid Tree

Description:

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.

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.

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.

Code:

class Solution:
    """
    @param n: An integer
    @param edges: a list of undirected edges
    @return: true if it's a valid tree, or false
    """
    def validTree(self, n, edges):
        # write your code here
        root = [i for i in range(n)]
        for i in edges:
            root1 = self.findRoot(root, i[0])
            root2 = self.findRoot(root, i[1])
            if root1==root2:
                return False
            else:
                root[root2] = root1
        return len(edges)==n-1

    def findRoot(self, root, n):
        if root[n]==n:
            return n
        else:
            return self.findRoot(root, root[n])