Functional, but not recursive, tree traversal

Posted on

Problem

To learn tree traversal i implemented os.walk, tail-recursive and stack-based. Unfortunately Python doesn’t support tail call optimization. How do i rewrite my function so that it is maximally functional paradigm compliant? I want to keep side-effects to minimum yet not dive into deep recursions.

def walk_rec(path):
    def i(result, stack):
        if not stack:
            return result
        else:
            path = stack.pop()
            with cd(path):
                ds, fs = partition(isdir, listdir(path))
            stack.extend(join(path, d) for d in ds)
            result.append((path, ds, fs))
            return i(result, stack)
    return i([], [expanduser(path)])

def walk_iter(path):
    stack = [expanduser(path)]
    result = []
    while stack:
        path = stack.pop()
        with cd(path):
            ds, fs = partition(isdir, listdir(path))
        stack.extend(join(path, d) for d in ds)
        result.append((path, ds, fs))
    return result

See: cd, partition.

I could get rid of cd, but this is not critical.

with cd(path):
    ds, fs = partition(isdir, listdir(path))
# can be replaced with
ds, fs = partition(lambda d: isdir(join(path, d)),
                   listdir(path))

Solution

Does this help? In your walk_iter, instead of using pop, grab the first element of the stack. Then do something like this …

if first_element_has_sub_elements:
    new_stack = stack[0][1]
    new_stack.extend(stack[1:])
    stack = new_stack[:]
else:
    stack = stack[1:]

This was my experiment, in case it’s helpful …

TREE = [
 ("L1A", [("L1A2A",), 
          ("L1A2B", [("L1A2B3A", )])]),
 ("L1B", [("L1B2A", [("L1B2A3A", )])]),
 ("L1C",),
 ("L1D", [("L1D2A",), 
          ("L1D2B", [("L1D2B3A", ),
                     ("L1D2B3B",)])]),
 ]

EXPECTED = ["L1A", "L1A2A", "L1A2B", "L1A2B3A",
            "L1B", "L1B2A", "L1B2A3A", "L1C", 
            "L1D", "L1D2A", "L1D2B", "L1D2B3A", "L1D2B3B"]

def traverse_tree(TREE):
    chopped_tree = TREE[:]
    results = []
    while chopped_tree:
        results.append(chopped_tree[0][0])
        if len(chopped_tree[0]) > 1:
            new_tree = chopped_tree[0][1]
            new_tree.extend(chopped_tree[1:])
            chopped_tree = new_tree[:]
        else:
            chopped_tree = chopped_tree[1:]

    return results

if __name__ == "__main__":
    results = traverse_tree(TREE)
    if results == EXPECTED:
        print "OK"
    else:
        print "MISMATCH"
    print "Expected: "
    print EXPECTED
    print "Results: "
    print results

Leave a Reply

Your email address will not be published. Required fields are marked *