Дейкстра с нуля

Я реализовал алгоритм Дейкстры с нуля и задаюсь вопросом, могу ли я сохранить код? Весь алгоритм здесь.

n — количество узлов, m — количество ребер. 1 ≤ 𝑛 ≤ 10 ^ 4, 0 ≤ 𝑚 ≤ 10 ^ 5, 𝑢 ̸ = 𝑣, 1 ≤ 𝑢, 𝑣 ≤ 𝑛, веса ребер — неотрицательные целые числа, не превышающие 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

И любые предложения по коду также были бы очень признательны. Заранее спасибо.

0

Добавить комментарий

Ваш адрес email не будет опубликован. Обязательные поля помечены *