zhipingchen / datastructure Goto Github PK

# graph node definition
class Node(object):
    def __init__(self, index, weight, next=None):
        self.index = index
        self.weight = weight
        self.next = next


# adjacency list definition
class AdjacencyList(object):
    def __init__(self, number):
        self.number = number
        self.list = [None] * number

    # insert node
    def insert(self, origin, index, weight=1):
        node = Node(index, weight, self.list[origin - 1])
        self.list[origin - 1] = node


# graph edge definition
class Edge(object):
    def __init__(self, begin, end, weight):
        self.begin = begin
        self.end = end
        self.weight = weight


def relax(edge, distance, parent):
    if distance[edge.begin - 1] is None:
        pass
    elif distance[edge.end - 1] is None or distance[edge.end - 1] > distance[edge.begin - 1] + edge.weight:
        distance[edge.end - 1] = distance[edge.begin - 1] + edge.weight
        parent[edge.end - 1] = edge.begin - 1
        return True
    return False


# get edges from adjacency list
def getEdgesFromAdjacencyList(graph):
    edges = []
    for i in range(graph.number):
        node = graph.list[i]
        while node:
            edge, node = Edge(i + 1, node.index, node.weight), node.next
            edges.append(edge)
    return edges


def bellman_ford(graph, start, end):
    distance, parent = [None for i in range(graph.number)], [None for i in range(graph.number)]
    times, distance[start - 1], edges = 1, 0, getEdgesFromAdjacencyList(graph)
    flag = True
    while flag and times < graph.number:
        flag = False
        for i in range(len(edges)):
            if relax(edges[i], distance, parent) and not flag:
                flag = True
        times += 1
    for i in range(len(edges)):
        if relax(edges[i], distance, parent):
            return False
    return True

if __name__ == '__main__':
    graph = AdjacencyList(4)
    graph.insert(1, 2, 4)
    graph.insert(1, 3, 6)
    graph.insert(1, 3, -10)
    graph.insert(2,3,2)
    graph.insert(3,4,5)
    print(bellman_ford(graph, 1, 4))