Chương trình tìm số bus tối thiểu đến đích cuối trong Python

Cho ma trận n×3, mỗi hàng [src, dest, id] = tuyến bus từ src đến dest. Lên bus mới tốn 1 đơn vị tiền, ở trên bus cũ chỉ tốn 1 đơn vị (dòng gốc: "takes one unit of money to get on a new bus, but if we stay on the same bus we have to pay one unit only" - hơi mơ hồ, hiểu là: chuyển bus tốn +1 chi phí, ở lại bus hiện tại tốn 0 thêm). Tìm chi phí tối thiểu từ vị trí 0 đến trạm cuối (largest location). Không có giải pháp → -1.

Ví dụ: [[0,1,0],[1,2,0],[2,3,0],[3,5,1],[5,0,2]] → output 2 (bus 0: 0→3, bus 1: 3→5).

Thuật toán Dijkstra trên trạng thái (vị trí, bus hiện tại):

• Xây đồ thị adj[src] = [(dest, bus_id)]
• Priority queue (min-heap) lưu (cost, pos, cur_bus)
• Khởi tạo: (0, 0, -1) (tại vị trí 0, chưa lên bus)
• Pop trạng thái chi phí nhỏ nhất:
  • Nếu pos == target → return cost
  • Nếu (pos, cur_bus) đã visited → continue
  • Đánh dấu visited
  • Duyệt cạnh (nex_pos, nex_bus) từ pos:
    • next_cost = cost + (1 if nex_bus != cur_bus else 0)
    • Push (next_cost, nex_pos, nex_bus)
• Hết heap → return -1

Triển khai

from collections import defaultdict
from heapq import heapify, heappop, heappush

class Solution:
    def solve(self, connections):
        start = 0
        target = max(max(y, x) for y, x, _ in connections)
        adj = defaultdict(list)
        for f, t, id in connections:
            adj[f].append((t, id))
        
        hp = [(0, start, -1)]
        seen = defaultdict(set)
        
        while hp:
            cost, cur_pos, cur_bus = heappop(hp)
            if cur_pos == target:
                return cost
            if cur_bus in seen[cur_pos]:
                continue
            seen[cur_pos].add(cur_bus)
            for nex_pos, nex_bus in adj[cur_pos]:
                next_cost = cost
                if nex_bus != cur_bus:
                    next_cost += 1
                heappush(hp, (next_cost, nex_pos, nex_bus))
        return -1

ob = Solution()
matrix = [
    [0, 1, 0], [1, 2, 0], [2, 3, 0],
    [3, 5, 1], [5, 0, 2]
]
print(ob.solve(matrix))

Đầu vào

[[0, 1, 0], [1, 2, 0], [2, 3, 0], [3, 5, 1], [5, 0, 2]]

Đầu ra

2

Giải thích: Trạng thái = (vị trí, bus_id). Chi phí = số lần chuyển bus. Dijkstra đảm bảo tìm chi phí nhỏ nhất. seen[pos] lưu set bus_id đã xử lý tại vị trí đó. Target = max location trong input.