Chương trình kiểm tra đường đi chi phí thấp hơn trong đồ thị vô hướng Python

Cho đồ thị vô hướng có trọng số. Cài đặt hàm query(p, q, limit): kiểm tra có đường đi từ p đến q với chi phí < limit không (hoặc ≤ limit tùy đề bài). Trả về True/False.

Ví dụ: Đồ thị 5 đỉnh, các cạnh: (0,1,10), (0,2,20), (1,4,10), (0,3,10), (1,2,20), (2,3,10). Queries: (0,2,10)→False, (3,1,30)→True, (4,3,30)→True.

Giải thuật (Offline + Union-Find + Sort)

• Sắp xếp cạnh theo trọng số tăng.
• Xử lý cạnh theo nhóm cùng trọng số: union-find kết nối các đỉnh.
• Lưu snapshot parent array sau mỗi nhóm trọng số vào connections, trọng số vào weights.
• Query(limit): binary search (bisect_left) trên weights tìm index. Nếu index=0 → False. Kiểm tra connections[index-1][p] == connections[index-1][q].

Code Python

import bisect

class Solution:
    def __init__(self, n, edgeList):
        def find(node):
            if parent[node] != node:
                parent[node] = find(parent[node])
            return parent[node]
        def union(x, y):
            parent[find(y)] = find(x)
        
        parent = {i: i for i in range(n)}
        edgeList.sort(key=lambda x: x[2])
        self.connections = []
        self.weights = []
        for idx, (i, j, w) in enumerate(edgeList):
            union(i, j)
            if idx != len(edgeList) - 1 and w == edgeList[idx + 1][2]:
                continue
            self.weights.append(w)
            self.connections.append([find(i) for i in parent])
    
    def query(self, p, q, limit):
        index = bisect.bisect_left(self.weights, limit)
        if index == 0:
            return False
        return self.connections[index - 1][p] == self.connections[index - 1][q]

ob = Solution(5, [[0,1,10],[0,2,20],[1,4,10],[0,3,10],[1,2,20],[2,3,10]])
print(ob.query(0, 2, 10))  # False
print(ob.query(3, 1, 30))  # True
print(ob.query(4, 3, 30))  # True

Đầu vào / Đầu ra

Output: False True True