โปรแกรมตรวจสอบเส้นทางที่มีข้อจำกัดความยาวในกราฟใน Python

มีกราฟไม่มีทิศทาง น้ำหนัก n โหนด edgeList[i]=(u,v,w) และคิวรี (p,q,lim) ถามว่ามีเส้นทางจาก p ไป q ระยะทาง < lim หรือไม่ คืน True/False แต่ละคิวรี

ตัวอย่าง: ผลลัพธ์ [True, False, True]

วิธีแก้ปัญหา (Union-Find + Offline)

1. เรียง edgeList ตามน้ำหนัก
2. เรียงคิวรีตาม lim
3. Union-Find เพิ่ม edge ทีละเส้นจนน้ำหนัก < limit ของคิวรีนั้น
4. ตรวจสอบว่า p,q อยู่ใน component เดียวกันหรือไม่

ตัวอย่างโค้ด

def solve(n, edgeList, queries):
    parent = [i for i in range(n+1)]
    rank = [0]*(n+1)
    def find(parent, x):
        if parent[x] == x: return x
        parent[x] = find(parent, parent[x])
        return parent[x]
    def union(parent, a, b):
        a = find(parent, a); b = find(parent, b)
        if a == b: return
        if rank[a] < rank[b]: parent[a] = b
        elif rank[a] > rank[b]: parent[b] = a
        else: parent[b] = a; rank[a] += 1
    edgeList.sort(key=lambda x: x[2])
    res = [0]*len(queries)
    queries = [[i, ch] for i, ch in enumerate(queries)]
    queries.sort(key=lambda x: x[1][2])
    ind = 0
    for i, (a, b, w) in queries:
        while ind < len(edgeList) and edgeList[ind][2] < w:
            union(parent, edgeList[ind][0], edgeList[ind][1])
            ind += 1
        res[i] = find(parent, a) == find(parent, b)
    return res

n = 4
edgeList = [(1,2,16),(1,3,8),(2,4,3),(2,3,6),(4,3,3)]
queries = [(1,4,12),(2,3,3),(1,2,15)]
print(solve(n, edgeList, queries))

อินพุต

4, [(1,2,16),(1,3,8),(2,4,3),(2,3,6),(4,3,3)], [(1,4,12),(2,3,3),(1,2,15)]

เอาต์พุต

[True, False, True]