cp_library/alg/graph/shortest_path_fn.py

• View this file on GitHub
• Last update: 2025-03-30 20:17:47+09:00

Depends on

• cp_library/ds/heap/heap_proto.py
• cp_library/ds/heap/min_heap_cls.py

Code

import cp_library.alg.graph.__header__
from math import inf
from cp_library.ds.heap.min_heap_cls import MinHeap

def shortest_path(G, s: int, g: int) -> tuple[list[int]|None,list[int]]:
    D = [inf] * G.N
    D[s] = 0
    if s == g:
        return [], D
    par = [-1] * G.N
    par_edge = [-1] * G.N
    Eid = G.edge_ids()
    heap = MinHeap()
    heap.push((0, s))

    while heap:
        d, v = heap.pop()
        if d > D[v]: continue
        if v == g: break
    
        for i,(u, w, *_) in enumerate(G[v]):
            if (nd := d + w) < D[u]:
                D[u] = nd
                par[u] = v
                par_edge[u] = Eid[v][i]
                heap.push((nd, u))
    
    if D[g] == inf:
        return None, D
        
    path = []
    current = g
    while current != s:
        path.append(par_edge[current])
        current = par[current]
        
    return path[::-1], D

'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
             https://kobejean.github.io/cp-library               
'''
from math import inf

from collections import UserList
from typing import Iterable
from heapq import heapify, heappop, heappush, heappushpop, heapreplace
from typing import Generic
from typing import TypeVar
_T = TypeVar('T')

class HeapProtocol(Generic[_T]):
    def pop(self) -> _T: ...
    def push(self, item: _T): ...
    def pushpop(self, item: _T) -> _T: ...
    def replace(self, item: _T) -> _T: ...

class MinHeap(HeapProtocol[_T], UserList[_T]):
    def __init__(self, iterable: Iterable = None):
        super().__init__(iterable)
        heapify(self.data)
    
    def pop(self): return heappop(self.data)
    def push(self, item: _T): heappush(self.data, item)
    def pushpop(self, item: _T): return heappushpop(self.data, item)
    def replace(self, item: _T): return heapreplace(self.data, item)

def shortest_path(G, s: int, g: int) -> tuple[list[int]|None,list[int]]:
    D = [inf] * G.N
    D[s] = 0
    if s == g:
        return [], D
    par = [-1] * G.N
    par_edge = [-1] * G.N
    Eid = G.edge_ids()
    heap = MinHeap()
    heap.push((0, s))

    while heap:
        d, v = heap.pop()
        if d > D[v]: continue
        if v == g: break
    
        for i,(u, w, *_) in enumerate(G[v]):
            if (nd := d + w) < D[u]:
                D[u] = nd
                par[u] = v
                par_edge[u] = Eid[v][i]
                heap.push((nd, u))
    
    if D[g] == inf:
        return None, D
        
    path = []
    current = g
    while current != s:
        path.append(par_edge[current])
        current = par[current]
        
    return path[::-1], D

Back to top page