Walking over a N branching tree

Question

The aim of this code, written in Python3.6, was to write a function that will non-recursively walk over a N-branching tree structure.

It works by using a list of indices that keep track of visited nodes and by retracing prior steps when a tree branch ends. The function non_recurs visits nodes top down left to right (as far as I can tell).

non-recurs allows the caller to provide a function to the func argument which will be applied to to each node. The end_only argument will determine if the supplied function will be applied to all nodes (end_only=False) or only nodes at the end of branches (end_only=True).

1st I would like to know if this attempt is actually capable of walking over N-branching trees and 2nd I would like to know what you think about my implementation.

---

NOTE: There are two sections to my code below. The first is to generate N-branching trees and is not the focus of this code review. But is required to generate my demonstrated output

---

EDIT

1. I added IndexError to the except of the try block

tree generation code

from itertools import product

# CODE to create test tree's for walking
def assign(tree, idx, value):
    new_idx = idx[:]
    assignment_idx = new_idx.pop()
    ref = tree
    for i in new_idx:
        ref = ref[i]
    ref[assignment_idx] = value
    return tree

def n_branch_tree(height, branches):
    idx = ([i] for i in range(branches))
    tree = list(range(branches))
    count = 0
    node = 0
    while count < height:
        for i in idx:
            # mutates tree
            assign(tree, list(i), list(range(node, node+branches)))
            node += branches
        count += 1
        idx = product(range(branches), repeat=count)
    return tree

---

tree walk code

# Code to walk over tree

def walk(tree, idx):
    """Return tree node at provided index

    args:
        tree: tree to index
        idx: index of desired node

    returns: node
    """
    for i in idx:
        tree = tree[i]
    return tree

def non_recurs(tree, func=print, branches=2, end_only=True, ):
    """Non-recursively walk n-branching tree

    args:
        tree: n-branching tree
        func: function that takes tree node as first argument
        branches: The number of branches each node has
        end_only: Default is True. When  True will only apply func to
                end nodes. When False will apply func to all Nodes

    """
    branches = branches - 1  # Because index's start at 0
    idx = [0]
    node = None
    while True:
        # print(idx)
        try:
            node = walk(tree, idx)
        except (TypeError, IndexError):
            # Could not find node at index
            try:
                # Walk back up tree until a node has not
                # been fully explored
                while idx[-1] == branches:
                    idx.pop()
            except IndexError:
                # Means all nodes in index have been fully explored
                break

            # Increase index if current index is under node branch limit
            if idx[-1] < branches:
                idx[-1] += 1

            # Apply func to end node
            if end_only and node:
                func(node)
                node = None
        else:
            idx.append(0)

            # Apply func to all nodes
            if not end_only:
                func(node)

if __name__ == '__main__':

        tree = n_branch_tree(height=3, branches=3)

        print(tree)

        value_store = []

        non_recurs(tree, func=value_store.append, branches=3, end_only=True)

        print(value_store)

---

Outputs

[[[18, 19, 20], [21, 22, 23], [24, 25, 26]], [[27, 28, 29], [30, 31, 32], [33, 34, 35]], [[36, 37, 38], [39, 40, 41], [42, 43, 44]]]
[18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44]

Answer 1

About your solution

Your tree generation code is wrong. It doesn't return a tree, it returns a multidimensional list instead. There is a good, detailed article about a tree data structure: Everything you need to know about tree data structures.

The main idea of tree - the connection of nodes. All begins from the root - the singular node at the first level, that contains references to the all nodes at the next level, which in their turn have references to the next level nodes, so on. Thus, every node should have as minimum as two fields: one for value and one for list of its childs (terminology used in trees).

Edit after the comment [start]:

Actually, the tree generation code returned value is resembling the tree, especially taking into account that the Python lists and numbers are objects. It has root object with 3 childs, each child has three childs too, so on. (when height=3, branches=3). At the last level it has objects with a number value.

Thus, it complies the first tree data structure requirement: the connection of nodes (from my own definition :)). But not all nodes have value - only the last level's nodes do (the leaves in the tree terminology). So, you anyway can't walking through tree, changing or printing all nodes value, because some nodes don't have them:

root[
    2-nd lvl[         4-th lvl with values
            3-rd lvl[18, 19, 20],
                    [21, 22, 23],
                    [24, 25, 26] 
            ],  
            [   
                    [27, 28, 29],
                    [30, 31, 32],
                    [33, 34, 35] 
            ],  
            [   
                    [36, 37, 38],
                    [39, 40, 41],
                    [42, 43, 44] 
            ]   
    ]  

Edit after the comment [end]:

I didn't throughout investigation of your tree walk code, but by glance it uses the the wrong tree idea as well as the tree generation code and therefore, can't work correctly.

My partial solution

I read your requirements and wrote my own solution, that creates tree, walks through it and applies a passed function to each node. It can't process only end nodes though, but this functionality is easy to add.

class Tree(object):

    def __init__(self):
        self.value = None
        self.childs = None

    # Build tree recursively
    # All nodes doesn't have values, just child list
    # Values will be set by the passed function
    def grow_tree(self, height, child_num):
        if height < 2:
            return

        self.childs = []    

        for i in range(child_num):
            child = Tree()
            self.childs.append(child)
            child.grow_tree(height - 1, child_num)

    # Walk through tree iteratively
    def walk_tree(self, func):

        all_current_level_nodes = [self]

        while all_current_level_nodes:
            all_next_level_nodes = []

            for node in all_current_level_nodes:
                func(node)
                if node.childs:
                    all_next_level_nodes.extend(node.childs)    

            all_current_level_nodes = all_next_level_nodes


### Recursive implementation
###
#   def walk_tree(self, func):
#       func(self)
#
##      if isinstance(self.childs, list):
#       if self.childs: 
#           for child in self.childs:
#               child.walk_tree(func)       

Testing

### Functions for passing to the "walk_tree" method         
def set_values(node):
    node.value = set_values.cnt
    set_values.cnt += 1

def print_values(node):
    print(node.value)

def print_id(node):
    print(id(node))

def square_values(node):
    node.value = node.value ** 2

tree = Tree()

tree.grow_tree(2, 3)

tree.walk_tree(print_id)

# Function is an object. Add the "cnt" field to this object.
# It will work as a static variable, preserve counter between calls.
set_values.cnt = 1
tree.walk_tree(set_values)

tree.walk_tree(print_values)

tree.walk_tree(square_values)

tree.walk_tree(print_values)

Output

# "print_id" function was applied on each node while walk_tree
139632471400688
139632471400800
139632471400856
139632471400912
# "set_values" 
# then "print_values" were applied
1
2
3
4
# "square_values"
# then "print_values" were applied
1
4
9
16