class Graph:
    def dijkstra(self, graph, start):
        '''
        重要的是定义成本，可能是路径，可能是花费, 如果有限制次数，也可以将次数也算进成本里
        如果是次数
        :param graph:  存了每个点的next node
        '''
        import heapq
        heap = [(0, start)]  # cost from start node,end node
        dist = {}
        # self.path = [-1 for _ in range(len(graph.keys()))]
        while heap:
            (cost, u) = heapq.heappop(heap)
            if u in dist:
                continue
            dist[u] = cost
            for v, c in graph[u]:
                if v in dist:
                    continue
                next = cost + c
                # self.path[v] = u
                heapq.heappush(heap, (next, v))
        return dist

    def get_dijkstra_path(self, j):
        if self.path[j] == -1:
            return [j]
        else:
            i = self.path[j]
            return self.get_dijkstra_path(i) + [j]

    def floyd_warshall(self, grid):
        if not grid or not grid[0]:
            return
        n = len(grid)
        # self.path = [[-1] * n for _ in range(n)]
        # grid: n*n
        for k in range(n):
            for i in range(n):
                for j in range(n):
                    if grid[i][k] + grid[k][j] < grid[i][j]:
                        grid[i][j] = grid[i][k] + grid[k][j]
                        # self.path[i][j] = k
        return grid

    def get_floyd_warshall_path(self, i, j):
        if self.path[i][j] == -1:
            return [i, j]
        else:
            k = self.path[i][j]
            return self.get_floyd_warshall_path(i, k) + self.get_floyd_warshall_path(k, j)

    def bellman_ford(self, edges, n, start):
        """
        Bellman-Ford 算法计算单源最短路径
        :param edges: 所有边，形如 [(u, v, w)]
        :param n: 节点数
        :param start: 起点
        :return: 距离数组和路径数组
        """
        dist = [float('inf')] * n
        prev = [-1] * n
        dist[start] = 0

        for _ in range(n - 1):
            for u, v, w in edges:
                if dist[u] != float('inf') and dist[u] + w < dist[v]:
                    dist[v] = dist[u] + w
                    prev[v] = u

        # 检测负权环
        for u, v, w in edges:
            if dist[u] != float('inf') and dist[u] + w < dist[v]:
                raise ValueError("图中包含负权环")

        self.bf_path = prev
        return dist

    def get_bellman_ford_path(self, j):
        if self.bf_path[j] == -1:
            return [j]
        else:
            return self.get_bellman_ford_path(self.bf_path[j]) + [j]

    def dijkstra(self, graph, start, end):
        heap = [(0, start)]  # cost from start node,end node
        dist = {}
        while heap:
            (cost, u) = heapq.heappop(heap)
            if u in dist:
                continue
            dist[u] = cost
            for v, c in graph[u]:
                if v in dist:
                    continue
                next = cost + c
                heapq.heappush(heap, (next, v))
        return dist

# Bellman-Ford implementation from MIT 6006 course lesson #17
import math
import networkx as nx

# utility: Graph
class Graph:
    def __init__(self, vertices):
        self.V = vertices
        self.edges = []
    
    def add_edge(self, edge):
        self.edges.append(edge)


def bellmanFord(graph, s):

    d = dict.fromkeys(graph.V, math.inf) # distance pair 
                                         # will have default value of Infinity
    pi = dict.fromkeys(graph.V, None) # map of parent vertex
    
    # initialize
    d[s] = 0
    
    def relax(u, v, w):
        if d[v] > d[u] + w:
            d[v] = d[u] + w
            pi[v] = u
    
    # The length of the longest simple path(path w/o cycle) would be |V| - 1.
    # For example, you need 2 edges to connect 3 vertices.
    # Otherwise, there exists a negative cycle.
    for i in graph.V[:-1]:
        for u, v, w in graph.edges:
            relax(u, v, w)
            
    for u, v, w in graph.edges:
        # even after relaxing all the edges for |V| - 1 times,
        # we still have the posibillity to improve the existing path
        # this means there are negative cycles
        if d[v] > d[u] + w:
            return f'there exists a negetive cycle!'
                
    return d, pi

def shortest_path(s, t):
    try:
        d, pi = bellmanFord(g, s)
    except ValueError:
        return 'you can\'t find shortest path if the graph has negative cycle!'
    
    path = [t]
    current = t
    
    # if parent pointer is None,
    # then it's the source vertex
    while pi[current]:
        path.insert(0, pi[current])
        # set current to parent
        current = pi[current]
    
    if s not in path:
        return f'unable to find shortest path staring from "{s}" to "{t}"'
    
    return f'{" > ".join(path)}'

g = Graph(['A', 'B', 'C', 'D', 'E'])

# graph with negative cycle
nc_edges = [('A', 'B', 5), ('B', 'C', -1), ('C', 'D', 2), ('D', 'B', -2), ('C', 'E', 4)]

# w/o negative cycles
edges = [\
    ('A', 'B', 10), ('A', 'C', 3), ('B', 'C', 1), ('C', 'B', 4), \
    ('B', 'D', 2), ('C', 'D', 8), ('D', 'E', 7), ('E', 'D', 9), ('C', 'E', 2)]

# used for both finding shortest path and drawing graph
current_edge_group = edges

for edge in current_edge_group:
    g.add_edge(edge)

print( shortest_path('A', 'E') )

G = nx.DiGraph()
G.add_weighted_edges_from(current_edge_group)
nx.draw(G, with_labels = True, node_color='b', font_color='w')

class Solution:
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, K: int) -> int:
        graph = collections.defaultdict(list)
        for u, v, w in flights:
            graph[u].append((v, w))
        visit = {}
        hq = [(0, src, 0)]
        while hq:
            cost, u, steps = heapq.heappop(hq)
            if (u, steps) in visit:
                continue
            visit[(u, steps)] = 1
            if steps > K + 1:
                continue
            if u == dst:
                return cost
            for v, w in graph[u]:
                if (v, steps+1) not in visit:
                    heapq.heappush(hq, (cost + w, v, steps + 1))
        return -1

class Solution:
    def findCheapestPrice(self, n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
        inf = float("inf")
        g = [inf for _ in range(n)]
        g[src] = 0
        ans = inf
        # 正是belman-ford思想, 1次relax = 1次中转.由于我们最多只能中转 k 次，也就是最多搭乘 k+1 次航班
        for i in range(k+1):
            f = [inf for _ in range(n)]
            for u, v, w in flights:
                f[v] = min(f[v], g[u] + w)
            g = f
            ans = min(ans, f[dst])
        return ans if ans < inf else -1