HiveBrain v1.2.0
Get Started
← Back to all entries
patternpythonMinor

Breadth-first search on a tree in Python

Submitted by: @import:stackexchange-codereview··
0
Viewed 0 times
searchbreadthfirstpythontree

Problem

The following code performs breadth first search on a mockup tree

# some tree
#             A            - LAYER 0
#         B       C        - LAYER 1 
#      D     E             - LAYER 2
#   F                      - LAYER 3

def get_something(x):
    if x == 'a':
        return ['b', 'c']
    if x == 'b':
        return ['d', 'e']
    if x == 'c':
        return []
    if x == 'd':
        return ['f']
    if x == 'e':
        return []
    if x == 'f':
        return []


and it is traversed by the following not so beautiful code

# I mean it works, but looking at it makes me feel sick

def recursive_get(x, loop_limit, count_limit):

    # stop condition 1: loop limit exceeded and nonzero
    # stop condition 2: count limit exceeded and nonzero
    # stop condition 3: we have reached the bottom

    current = 0
    seen_list = []
    todo_list = [x]

    while ((current < loop_limit or loop_limit == 0) and
           (len(seen_list) < count_limit or count_limit == 0) and
           todo_list):

        result = get_something(todo_list.pop(0))
        seen_list.extend(result)
        todo_list.extend(result)
        current += 1

    return seen_list


Is there any syntactic sugar to get the second part to look more sexy? (less lines, syntactically nicer looking)?

Solution

get_something is not a good name for the function. Obviously, it specifies your tree by returning the child nodes of an argument node; rename to, say, get_neighbors. Also, your implementation overkills breadth-first search a bit. See what I mean:

def get_neighbors(x):
    if x == 'a':
        return ['b', 'c']
    if x == 'b':
        return ['d', 'e']
    if x == 'c':
        return []
    if x == 'd':
        return ['f']
    if x == 'e':
        return []
    if x == 'f':
        return []

    return [] # Other nodes have no neighbors.

def bfs(start):
    queue = [start]
    visited = set(start)

    while queue:
        node = queue.pop(0)

        for neighbor in get_neighbors(node):
            if neighbor not in visited:
                queue.append(neighbor)
                visited.add(neighbor)

    return visited

print(str(bfs('a')))


Also, note that seen_list should be rather a Python set as operating on it is more efficient than on lists.

Code Snippets

def get_neighbors(x):
    if x == 'a':
        return ['b', 'c']
    if x == 'b':
        return ['d', 'e']
    if x == 'c':
        return []
    if x == 'd':
        return ['f']
    if x == 'e':
        return []
    if x == 'f':
        return []

    return [] # Other nodes have no neighbors.

def bfs(start):
    queue = [start]
    visited = set(start)

    while queue:
        node = queue.pop(0)

        for neighbor in get_neighbors(node):
            if neighbor not in visited:
                queue.append(neighbor)
                visited.add(neighbor)

    return visited

print(str(bfs('a')))

Context

StackExchange Code Review Q#123534, answer score: 4

Revisions (0)

No revisions yet.