โปรแกรมหาเส้นทางที่มี penalty น้อยที่สุดระหว่างสองจุดยอดใน Python

สมมติว่าเราได้รับกราฟไม่มีทิศทางมีน้ำหนัก และต้องการหาเส้นทางที่มี penalty น้อยที่สุดจากโหนด a ไป b penalty ของเส้นทางคือ bitwise OR ของน้ำหนักเส้นขอบทั้งหมดในเส้นทาง เราต้องหาเส้นทาง "minimum penalty" นี้ และถ้าไม่มีเส้นทางระหว่างสองโหนด ให้ return -1

เช่น ถ้าอินพุตเป็นดังภาพ start (s) = 1, end (e) = 3 ผลลัพธ์จะเป็น 15 มีสองเส้นทางระหว่างจุดยอด 1 และ 3 เส้นทางเหมาะสมคือ 1→2→3 penalty คือ (10 OR 5) = 15

วิธีแก้ปัญหา

• กำหนดฟังก์ชัน helper() รับ G, s, e:
  • v := เซตใหม่
  • c := รายการใหม่ขนาด n เริ่มต้น infinity
  • heap := heap ใหม่มีคู่ (0, s)
  • while ขนาดของ heap > 0:
    • cst, cur := pop องค์ประกอบเล็กที่สุดจาก heap
    • c[cur] := min(cst, c[cur])
    • ถ้า (cst, cur) อยู่ใน v: ข้ามรอบถัดไป
    • ถ้า cur เท่ากับ e: return c[cur]
    • เพิ่ม (cst, cur) ลง v
    • สำหรับแต่ละ neighbor, n_cost ใน G[cur]: push ((n_cost OR cst), neighbor) ลง heap
  • return c[e]
• G := [รายการใหม่ที่มี n + 1 รายการว่าง]
• สำหรับแต่ละ item ใน edges: สร้าง G ด้วย (v, w) และ (u, w)
• ans := helper(G, s, e)
• return -1 ถ้า ans เท่ากับ inf มิฉะนั้น return ans

ตัวอย่าง

import heapq
from math import inf

def helper(G, s, e):
    v = set()
    c = [inf] * len(G)
    heap = [(0, s)]
    while len(heap) > 0:
        cst, cur = heapq.heappop(heap)
        c[cur] = min(cst, c[cur])
        if (cst, cur) in v:
            continue
        if cur == e:
            return c[cur]
        v.add((cst, cur))
        for neighbor, n_cost in G[cur]:
            heapq.heappush(heap, (n_cost | cst, neighbor))
    return c[e]

def solve(n, edges, s, e):
    G = [[] for _ in range(n + 1)]
    for item in edges:
        u, v, w = map(int, item)
        G[u].append((v, w))
        G[v].append((u, w))
    ans = helper(G, s, e)
    return -1 if ans == inf else ans

print(solve(4, [(1, 2, 10), (2, 3, 5), (2, 4, 15), (1, 4, 20)], 1, 3))

อินพุต

4, [(1, 2, 10), (2, 3, 5), (2, 4, 15), (1, 4, 20)], 1, 3

เอาต์พุต

15