network_delay.py

class Solution:
    def networkDelayTime(self, edges, N, K):
        """
        :type times: List[List[int]]
        :type N: int
        :type K: int
        :rtype: int
        """
        edges.sort(key=lambda e: e[2])
        # print(edges)
        (src_map, dst_map) = Solution.makeMaps(edges)
        src2 = src_map.copy()
        dst2 = dst_map.copy()
        src3 = src_map.copy()
        dst3 = dst_map.copy()
        self.edge_map = Solution.makeEdgeMap(edges)
        self.times = {K: 0}
        self.src_stack = []
        self.findTimes(0, src_map, K)

        if len(self.times) < N:
            return -1

        # self.reverseTimes(dst_map)
        self.convergeTimes(edges)

        print(self.times)
        return max(self.times.values())

    def findTimes(self, time_to, src_map, src):
        if src not in src_map:
            return
        dsts = src_map.pop(src)
        # print("findTimes: {}, {}".format(src, dsts))
        for (t, dst) in dsts:
            new_time_to = time_to + t
            self.setTime(dst, new_time_to)
            
        for (t, dst) in dsts:
            self.findTimes(time_to + t, src_map, dst)
            
        self.src_stack.append(src)
            
    def setTime(self, dst, time_to):
        if dst not in self.times:
            # print("initial time for {}: {}".format(dst, time_to))
            self.times[dst] = time_to
            return
        old_time = self.times[dst]
        if time_to < old_time:
            # print("faster time for {}: {}".format(dst, time_to))
            self.times[dst] = time_to
            
    def reverseTimes(self, dst_map):
        for dst in self.src_stack:
            if dst not in dst_map:
                continue
            srcs = dst_map.pop(dst)
            # print("reverseTimes: {} {}".format(dst, srcs))
            for (t, src) in srcs:
                self.setReverseTime(src, dst, t)
                
    def setReverseTime(self, src, dst, t):
        old_src = self.times[src]
        old_dst = self.times[dst]
        new_dst = old_src + t
        self.times[dst] = min(old_dst, new_dst)

    def convergeTimes(self, edges):
        loops = 0
        n_divergent = 1
        while n_divergent > 0:
            n_divergent = 0
            for e in edges:
                if self.convergeOneTime(e[0], e[1], e[2]):
                    n_divergent = n_divergent + 1
            # print("n_divergent: {}".format(n_divergent))
            loops = loops + 1
            print("convergence loops: {}".format(loops))
    
    def convergeOneTime(self, src, dst, t):
        old_src = self.times[src]
        old_dst = self.times[dst]
        new_dst = old_src + t
        if new_dst < old_dst:
            self.times[dst] = new_dst
            return True
        return False
        
    def makeMaps(edges):
        srcs = {}
        dsts = {}
        for e in edges:
            src = e[0]
            dst = e[1]
            time = e[2]
            if src in srcs:
                srcs[src].append((time, dst))
            else:
                srcs[src] = [(time, dst)]
            if dst in dsts:
                dsts[dst].append((time, src))
            else:
                dsts[dst] = [(time, src)]
        for src, dst_times in srcs.items():
            dst_times.sort()
        for dst, src_times in dsts.items():
            src_times.sort(reverse=True)
        return (srcs, dsts)
 
    def makeEdgeMap(edges):
        em = {}
        for e in edges:
            em[(e[0], e[1])] = e[2]
        return em