"""
p64 — Cheapest Flights Within K Stops (LeetCode 787, MEDIUM).

K stops = at most K+1 edges. Dijkstra fails (greedy cuts away viable
fewer-edge paths). Use Bellman-Ford with K+1 iterations OR BFS-with-state.
"""
from __future__ import annotations
import heapq
import random
from collections import defaultdict
from typing import List


INF = float("inf")


def find_cheapest_price_bf(n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
    prev = [INF] * n
    prev[src] = 0
    for _ in range(k + 1):
        curr = prev.copy()   # CRITICAL: copy so we relax exactly one more edge per iteration
        for u, v, w in flights:
            if prev[u] + w < curr[v]:
                curr[v] = prev[u] + w
        prev = curr
    return int(prev[dst]) if prev[dst] != INF else -1


def find_cheapest_price_bfs(n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
    adj: defaultdict = defaultdict(list)
    for u, v, w in flights:
        adj[u].append((v, w))
    # heap entries: (cost, node, stops_used)  -- stops_used = number of edges taken
    heap = [(0, src, 0)]
    # best[(node, stops_used)] to avoid re-expansion of dominated states
    best = {}
    while heap:
        cost, u, stops = heapq.heappop(heap)
        if u == dst:
            return cost
        if stops > k:
            continue
        if best.get((u, stops), INF) < cost:
            continue
        for v, w in adj[u]:
            new_cost = cost + w
            key = (v, stops + 1)
            if new_cost < best.get(key, INF):
                best[key] = new_cost
                heapq.heappush(heap, (new_cost, v, stops + 1))
    return -1


def find_cheapest_price_brute(n: int, flights: List[List[int]], src: int, dst: int, k: int) -> int:
    adj: defaultdict = defaultdict(list)
    for u, v, w in flights:
        adj[u].append((v, w))
    best = INF

    def dfs(u: int, cost: int, edges_used: int) -> None:
        nonlocal best
        if cost >= best:
            return
        if u == dst:
            best = cost
            return
        if edges_used == k + 1:
            return
        for v, w in adj[u]:
            dfs(v, cost + w, edges_used + 1)

    dfs(src, 0, 0)
    return int(best) if best != INF else -1


def edge_cases() -> None:
    assert find_cheapest_price_bf(3, [[0, 1, 100], [1, 2, 100], [0, 2, 500]], 0, 2, 1) == 200
    assert find_cheapest_price_bf(3, [[0, 1, 100], [1, 2, 100], [0, 2, 500]], 0, 2, 0) == 500
    assert find_cheapest_price_bfs(3, [[0, 1, 100], [1, 2, 100], [0, 2, 500]], 0, 2, 1) == 200
    assert find_cheapest_price_bfs(3, [[0, 1, 100], [1, 2, 100], [0, 2, 500]], 0, 2, 0) == 500
    # The classic counterexample for naive Dijkstra:
    flights = [[0, 1, 1], [0, 2, 5], [1, 2, 1], [2, 3, 1]]
    assert find_cheapest_price_bf(4, flights, 0, 3, 1) == 6
    assert find_cheapest_price_bfs(4, flights, 0, 3, 1) == 6
    assert find_cheapest_price_bf(2, [], 0, 1, 5) == -1
    assert find_cheapest_price_bf(1, [], 0, 0, 0) == 0
    print("edge_cases: PASS")


def stress_test() -> None:
    rng = random.Random(42)
    for trial in range(300):
        n = rng.randint(2, 5)
        m = rng.randint(0, n * (n - 1))
        flights = []
        for _ in range(m):
            u = rng.randrange(n)
            v = rng.randrange(n)
            if u != v:
                w = rng.randint(1, 9)
                flights.append([u, v, w])
        src = rng.randrange(n)
        dst = rng.randrange(n)
        k = rng.randint(0, n)
        a = find_cheapest_price_bf(n, flights, src, dst, k)
        b = find_cheapest_price_bfs(n, flights, src, dst, k)
        c = find_cheapest_price_brute(n, flights, src, dst, k)
        assert a == b == c, f"trial {trial}: n={n} flights={flights} src={src} dst={dst} k={k} bf={a} bfs={b} brute={c}"
    print("stress_test: 300 trials — Bellman-Ford, BFS-with-state, brute all agree.")


if __name__ == "__main__":
    edge_cases()
    stress_test()
    print("ALL TESTS PASSED")