I implemented the Dijkstra algorithm from scrath and wonder if I can save some code? The whole algorithm is here.
n is the number of nodes, and m the number of edges. 1 ≤ 𝑛 ≤ 10^4, 0 ≤ 𝑚 ≤ 10^5, 𝑢 ̸= 𝑣, 1 ≤ 𝑢, 𝑣 ≤ 𝑛, edge weights are non-negative integers not exceeding 10^8.
def restore_heap(heap_arr, dist, updated, vertices2indices):
"""
Restore the heap due to the updated nodes
@param heap_arr: a list, the heap with each element as vertice/node name
@param dist: a, list, the distance/priority of each node
@param updated: a list of tuples, updated nodes and the new distance in the previous relaxation
@param vertices2indices: a list, via which we can find the index of a node in the heap
"""
def restore(v):
# since the distance/priority of any updated node is smaller than before the updating, we
# only compare it with its parents
v_idx = vertices2indices[v]
# if we reach the root we stop the restoration
# assert v_idx >= 0, [heap_arr, dist, updated, v_idx]
if not v_idx:
return
parent_idx = (v_idx + 1) // 2 - 1
parent = heap_arr[parent_idx]
# if the priority of the updated node is smaller than its parent we swap the two
if dist[v] < dist[parent]:
heap_arr[v_idx], heap_arr[parent_idx] = heap_arr[parent_idx], heap_arr[v_idx]
# since the heap has been updated we should update the node-to-its-
# heap-index mapping
vertices2indices[v], vertices2indices[parent] = parent_idx, v_idx
restore(v)
else:
# if the priority of its parent is smaller than v we stop the restoration
return
for v, w in updated:
# important
dist[v] = w
restore(v)
def heap_pop(heap_arr, dist, vertices2indices):
"""
Pop the min of the heap(the first element) and restore the heap
@param heap_arr: a list, the heap
@param dist: a, list, the distance(or priority) of a node with the index as its node key
@param vertices2indices: a list, via which we can find the index of a node in the heap
return m: the value of the minimum node in the heap
"""
def restore(v):
# we compare the priority of the temperary root with its descendants
v_idx = vertices2indices[v]
# v has no children, we stop
if v_idx + 1 > len(heap_arr) // 2:
return
# (v_idx + 1) * 2 - 1
left_child_idx = 2 * v_idx + 1
# (v_idx + 1) * 2 + 1 - 1
right_child_idx = 2 * v_idx + 2
left_child = heap_arr[left_child_idx]
if right_child_idx > len(heap_arr) - 1:
right_child = left_child
else:
right_child = heap_arr[right_child_idx]
# choose the child with min priority/distance
min_child, min_child_idx = left_child, left_child_idx
if dist[left_child] > dist[right_child]:
min_child, min_child_idx = right_child, right_child_idx
# if the priority of the parent is greater than its min child in distance
if dist[v] > dist[min_child]:
heap_arr[v_idx], heap_arr[min_child_idx] = heap_arr[min_child_idx], heap_arr[v_idx]
# since the heap has been updated we should update the
# node-to-its-heap-index mapping
vertices2indices[v], vertices2indices[min_child] = min_child_idx, v_idx
restore(v)
else:
# if the priority of the parent is greater than v we stop the restoration
return
if len(heap_arr) == 1:
return heap_arr.pop()
m = heap_arr[0]
heap_arr[0] = heap_arr.pop()
vertices2indices[heap_arr[0]] = 0
restore(heap_arr[0])
return m
def heapify(heap_arr, s, vertices2indices):
"""
Move the source to the root of the heap
@param heap_arr: the heap
@param s: the source node
@param vertices2indices: a list, via which we can find the index of a node in the heap
"""
heap_arr[s], heap_arr[0] = heap_arr[0], heap_arr[s]
vertices2indices[s], vertices2indices[0] = 0, s
def distance(adj, cost, s, t):
"""
Given an directed graph with positive edge weights and with n vertices and
m edges as well as two vertices u and v, compute the weight of a shortest
path between u and v (that is, the minimum total weight of a path from u to v).
@param adj: int, represents the edges
@param cost: list of list, index represent nodes, and values the corresponding weights of edges
@param s: the source node
@param t: the destination node
@return: shortest distance from s to t
>>> distance([[1, 2], [2], [], [0]], [[1, 5], [2], [], [2]], 0, 2)
3
>>> distance([[1, 2], [2, 3, 4], [1, 4, 3], [], [3]], [[4, 2], [2, 2, 3], [1, 4, 4], [], [1]], 0, 4)
6
>>> distance([[1, 2], [2], []], [[7, 5], [2], []], 2, 1)
-1
>>> distance([[]], [[]], 0, 0)
0
>>> distance([[2, 3], [2, 4], [4, 3], [0, 4, 2], [1]], [[88526216, 28703032], [78072041, 22632987], [49438153, 12409612], [24629785, 96300300, 317823], [6876525]], 3, 1)
56632501
"""
# print(adj, cost, s, t)
heap_arr = list(range(len(adj)))
# map the node to its index in the heap arr
vertices2indices = list(range(len(adj)))
dist = [MAX_DISTANCE] * len(adj)
dist[s] = 0
heapify(heap_arr, s, vertices2indices)
while(heap_arr):
min_vertice = heap_pop(heap_arr, dist, vertices2indices)
if min_vertice == t and dist[min_vertice] < MAX_DISTANCE:
return dist[min_vertice]
updated = []
for v, w in zip(adj[min_vertice], cost[min_vertice]):
update = dist[min_vertice] + w
if dist[v] > update:
# should not update in one go, otherwise the unrestored nodes may
# block their descendants
# dist[v] = update
updated.append((v, update))
dist_c = [dist[i] for i in heap_arr]
restore_heap(heap_arr, dist, updated, vertices2indices)
return -1
And any suggestions for the code would also be very appreciated. Thanks in advance.
