Search code examples
pythonalgorithmtreepreorder

Preorder tree walk of a minimum spanning tree generated by Prim's algorithm


I'm trying to implement an approximate algorithm to solve the traveling salesman problem (TSP), which can be used when the triangle inequality holds for the edge weights. As described in Cormen et al., Introduction to Algorithms (3rd 3d.), the pseudocode is:

enter image description here

and here is an example:

enter image description here

What I'm struggling with is how to implement the preorder tree walk on the tree generated by Prim's algorithm. Since this is not a binary search tree, the pseudocode given at https://en.wikipedia.org/wiki/Tree_traversal#Pre-order_2 seems not to apply.

Rather, instead of having left and right attributes, the nodes have key and parent attributes. Here is how they are generated in my implementation of Prim's algorithm (with a small test case):

import math
import copy
import pytest
import pandas as pd
from cached_property import cached_property


class Node(object):
    def __init__(self, key=math.inf, parent=None):
        self.key = key
        self.parent = parent

    def __lt__(self, other):
        return self.key < other.key


class Graph(object):
    def __init__(self, edges):
        self.edges = edges

    @cached_property
    def nodes(self):
        _nodes = set()
        for edge in self.edges:
            _nodes.add(edge[0])
            _nodes.add(edge[1])
        return {node: Node() for node in list(_nodes)}

    @cached_property
    def adj(self):
        A = {node: [] for node in self.nodes}
        for edge in self.edges:
            u, v, _ = edge
            A[u].append(v)
            A[v].append(u)
        return A

    @cached_property
    def w(self):
        N = len(self.nodes)
        none_array = [[None for _ in range(N)] for _ in range(N)]
        df = pd.DataFrame(none_array, index=sorted(self.nodes), columns=sorted(self.nodes))
        for edge in self.edges:
            u, v, weight = edge
            df.set_value(u, v, weight)
            df.set_value(v, u, weight)
        return df

    def mst_prim(self, root):
        r = self.nodes[root]
        r.key = 0
        Q = copy.copy(self.nodes)
        while Q:
            u = min(Q, key=Q.get)
            u_node = Q.pop(u)
            for v in self.adj[u]:
                if v in Q and self.w[u][v] < self.nodes[v].key:
                    self.nodes[v].parent = u
                    self.nodes[v].key = self.w[u][v]


@pytest.fixture
def edges_simple():
    return [('a', 'b', 4),
            ('a', 'h', 8),
            ('b', 'h', 11),
            ('h', 'i', 7),
            ('b', 'c', 8),
            ('h', 'g', 1),
            ('i', 'c', 2),
            ('i', 'g', 6),
            ('c', 'd', 7),
            ('g', 'f', 2),
            ('c', 'f', 4),
            ('d', 'f', 14),
            ('d', 'e', 9),
            ('f', 'e', 10)
            ]

def test_mst_prim(edges_simple):
    graph = Graph(edges_simple)
    graph.mst_prim(root='a')
    # print("\n")
    # for u, node in graph.nodes.items():
    #   print(u, node.__dict__)
    assert graph.nodes['a'].parent is None
    assert graph.nodes['i'].parent == 'c'
    assert graph.nodes['d'].parent == 'c'



if __name__ == "__main__":
    # pytest.main([__file__+"::test_mst_prim", "-s"])
    pytest.main([__file__, "-s"])

How could I perform preorder tree traversal on this graph? (Note that this question is similar to pre-order traversal of a Minimum spanning tree, but I found the answer given there rather high-level).


Solution

  • I suggest you to add a new list in your Node class named children for example.

    After your Prim's algorithm you can run through your obtained nodes and add them to their parent's children. The complexity is O(n), so not a big deal. After that the DFS traversal will be easy.

    But again, as in the post you mentioned, you have to pick a order for your children for a preorder traversal. In your case when you only have reference to your parent, there's no way of knowing what is the left-most child for example.