Prove that the graph is a valid tree in Python

I recently solved this leetcode problem:


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.

For 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.



Note: 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.

My code is correct, and runtime beats 80% of answers submitted, but the overall structure is not at all user friendly. Could anyone give some code review to improve readability?

Don't worry what alternative method or improving the solution since it is already in \$O(n)\$.

`from collections import defaultdict

class Solution(object):
def validTree(self, n, edges):
"""
:type n: int
:type edges: List[List[int]]
:rtype: bool
"""
if not edges:
if n == 1:
return True
elif n == 2:
return False

visited_nodes = defaultdict(list)
count = True
unvisited = []
for key, value in edges:

if key in visited_nodes and value in visited_nodes:
return False

if count:
visited_nodes[key] = []
visited_nodes[value] = []
count = False
continue
if key in visited_nodes:
visited_nodes[key].append(value)
visited_nodes[value] = []
elif value in visited_nodes:
visited_nodes[value].append(key)
visited_nodes[key] = []
else:
unvisited.append([key, value])

for key, value in unvisited:

if key in visited_nodes and value in visited_nodes:
return Fals

Before going in depth of your algorithm, a few stylistics notes:

• I assume that the class Solution(object) thing is required by leetcode, so it might not apply. But as a general note, your validTree method does not need to store/interact with any state. Thus there is no need for it to be a method and would benefit to be defined as a function. Also the naming should be snake_case as per PEP8: def valid_tree(n, edges):. is_valid_tree might also be a better name.

• In Python 2, the keys method of dictionaries build and return a new list. It is generally recommended that your use iterkeys instead which return a proxy object to the actual keys of the dictionary. However, such object does not support the len protocol. But the dictionary support it without having to build anything new: len(visited_nodes).

Tidy up your special case

Your algorithm performs two checks:

• Every node is reached through an edge;

• Every node can reach each other through at most one path.

The only special case to that is, for only one node, to have no edges. The other cases will be handled by the for loops (which will do nothing if edges is empty) and the last check. Thus I would write:

def is_valid_tree(n, edges):
    if n == 1:
        # 1 node must not have edges
        return not edges
    ...

You can also remove all special cases by using default values in visited_nodes. More on that later.

Don't use stuff that you don't use

Sounds weird, doesn't it? But the same goes for your use of defaultdict: you use it without using the extra functionalities over dict. In fact, using visited_nodes = {} will not change anything in your program. So use that instead, it will remove a bit of overhead.

defaultdicts usage is to make use of default values instead of having to explicitly assign them, such as in:

d = defaultdict(list)
d['test'].append('something')
d['test'].append('something else')
d['other test'].append('a third thing')
print(d) # defaultdict(, {'test': ['something', 'something else'], 'other test': ['a third thing']})

But instead, here you explicitly assign empty lists before assigning to them. This is because you explicitly want to define the key of the dictionary; and, in fact, you do not care about the value. Since you do not use values, why using a dictionary to start with? Use a set they are like dictionaries, but only for keys:

def is_valid_tree(n, edges):
    if n == 1:
        return not edges

    visited_nodes = set()
    count = True
    unvisited = []

    for key, value in edges:
        if key in visited_nodes and value in visited_nodes:
            return False

        if count:
            visited_nodes.add(key)
            visited_nodes.add(value)
            count = False
            continue
        if key in visited_nodes:
            visited_nodes.add(value)
        elif value in visited_nodes:
            visited_nodes.add(key)
        else:
            unvisited.append([key, value])

    for key, value in unvisited:
        if key in visited_nodes and value in visited_nodes:
            return False

        if key in visited_nodes:
            visited_nodes.add(value)
        elif value in visited_nodes:
            visited_nodes.add(key)
        else:
            return False

    return n == len(visited_nodes)

Remove "duplicated" checks

First of, the count variable is useless: it only store whether the dictionary (or the set) is empty. You can get the exact same information using not visited_nodes which is True when visited_node is empty and False otherwise.

Second, you want to add both key and value as key (so maybe they need to be renamed) of visited_nodes if either of them is already one. You can thus simplify the writing using a boolean or and retain the same functionality:

if key in visited_nodes or value in visited_nodes:
    visited_nodes.add(key)
    visited_nodes.add(value)

Using the set approach also help you write less verbose code: use the intersection & operator of sets:

for edge in edges:
    edge = set(edge)
    commons = visited_nodes & edge

    # If a path exist between the nodes in edge
    if commons == edge:
        # A cycle exist and thus it is not a tree
        return False

    # If a node has been visited or if it is the first edge
    if commons or not visited_nodes:
        # Add the path
        visited_nodes.update(edge)
    else:
        # Process later
        unvisited.append(edge)

Third, as stated earlier, you can get rid of all your special cases by initializing visited_nodes = {0}. This will allow you to remove the first check (if n == 1) since it is now handled by return n == len(visited_nodes). And it will also handle adding branches from the node 0 instead of any two first nodes.

Last, the body of both for loops is pretty similar. You may want to factorize it out. Unfortunately, I'm not able to come up with something elegant. It is possible t

StackExchange Code Review Q#115776, answer score: 6