graph.py

class Node(object):
    def __init__(self, value):
        self.value = value
        self.edges = []


class Edge(object):
    def __init__(self, value, node_from, node_to):
        self.value = value
        self.node_from = node_from
        self.node_to = node_to


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

    def insert_node(self, new_node_val):
        new_node = Node(new_node_val)
        self.nodes.append(new_node)

    def insert_edge(self, new_edge_val, node_from_val, node_to_val):
        from_found = None
        to_found = None
        for node in self.nodes:
            if node_from_val == node.value:
                from_found = node
            if node_to_val == node.value:
                to_found = node
        if from_found is None:
            from_found = Node(node_from_val)
            self.nodes.append(from_found)
        if to_found is None:
            to_found = Node(node_to_val)
            self.nodes.append(to_found)
        new_edge = Edge(new_edge_val, from_found, to_found)
        from_found.edges.append(new_edge)
        to_found.edges.append(new_edge)
        self.edges.append(new_edge)

    def get_edge_list(self):
        # (Edge Value, From Node Value, To Node Value)
        edge_list = []
        for e in self.edges:
            edge_list.append((e.value, e.node_from.value, e.node_to.value))
        return edge_list

    def get_adjacency_list(self):
        # (To Node, Edge Value)
        max_val = self.get_max_vertex()
        a = [None for i in range(max_val+1)]
        for n in self.nodes:
            v = n.value
            q = []
            for e in n.edges:
                w = e.node_from.value
                if w == v:
                    q.append((e.node_to.value, e.value))
            if len(q) >= 1:
                a[v] = q
        return a

    def get_adjacency_matrix(self):
        # Row numbers represent from nodes,
        # column numbers represent to nodes.
        # Store the edge values in each spot,
        # and a 0 if no edge exists.
        a = self.get_adjacency_list()
        m = []
        max_val = self.get_max_vertex()
        for i in range(max_val+1):
            c = [0 for x in range(len(self.nodes)+1)]
            m.append(c)
        for j in range(len(a)):
            if a[j] is not None:
                for (v, e) in a[j]:
                    m[j][v] = e
        return m

    def get_max_vertex(self):
        max_val = 0
        for node in self.nodes:
            if node.value > max_val:
                max_val = node.value
        return max_val


graph = Graph()
graph.insert_edge(100, 1, 2)
graph.insert_edge(101, 1, 3)
graph.insert_edge(102, 1, 4)
graph.insert_edge(103, 3, 4)
# Should be [(100, 1, 2), (101, 1, 3), (102, 1, 4), (103, 3, 4)]
print(graph.get_edge_list())
# Should be [None, [(2, 100), (3, 101), (4, 102)], None, [(4, 103)], None]
print(graph.get_adjacency_list())
# Should be [[0, 0, 0, 0, 0], [0, 0, 100, 101, 102], [0, 0, 0, 0, 0], [0, 0, 0, 0, 103], [0, 0, 0, 0, 0]]
print(graph.get_adjacency_matrix())