Dijkstra with Heaps for problem set 5 udacity course intro to algorithms

dijkstra_with_heaps.py

# udacity ps 5.1
#


class BinaryHeap(object):

    """BinaryHeap for maintaining ordered
    list of nodes with smallest score at the
    top of the heap.

    Cannot have more than one node with same key."""

    def __init__(self):
        self.ordered = []
        self.length = 0
        self.map = {}

    def insert(self, key, score):
        """
        Append node to heap, update map and length,
        then upheapify on that node.
        """

        # update map
        i = self.length
        self.map[key] = i

        # append to heap
        self.ordered.append({"key": key, "score": score})

        # update length
        self.length += 1

        # no need to upheapify on one node
        if self.length > 1:
            self.upheapify(i)

    def pop(self):
        """
        Removes root node (smallest).
        Replaces this deleted root node with the last element from
        the ordered list and downheapifies.
        """

        # remove the root node from data structures
        smallest = self.ordered.pop(0)
        del self.map[smallest['key']]
        self.length -= 1

        if self.length == 0:

            return smallest

        elif self.length == 1:
            # already sorted
            self._update_map(0)

        elif self.ordered:
            # move last element to the root
            # and downheapify on that node
            self.ordered.insert(0, self.ordered.pop())
            self._update_map(0)
            self.downheapify(0)

        return smallest

    def decrement_score(self, key, new_score):
        """
        Get index of node from map.
        Update with new score.
        Upheapify on that node.
        """
        i = self.map[key]

        self.ordered[i]["score"] = new_score

        self.upheapify(i)

    def smallest_key(self):
        return self.ordered[0]["key"]

    def fetch_score(self, key):
        return self.ordered[self.map[key]]["score"]

    def upheapify(self, i):
        """
        If its not the root node and the parent is bigger
        swap it with the parent and iterate on its new position.
        """
        while i:

            # get parent
            current = self.ordered[i]
            j = self.parent(i)
            parent = self.ordered[j]

            if parent["score"] > current["score"]:

                # parent is greater swap
                self.swap(i, j)
                i = j

            else:

                # heap property maintained
                return

    def downheapify(self, i):
        """
        - if it is a leaf return
        - if it has one child and that child is smaller,
          swap it and finish.
        - if it has two children, swap it with the smallest one
          and iterate on its new position
        """

        while True:
            current = self.ordered[i]
            l = self.left_child(i)
            r = self.right_child(i)

            # if is a leaf
            if l >= len(self.ordered):
                return

            left = self.ordered[l]

            if self.one_child(i):

                if left["score"] < current["score"]:
                    # left child is smaller
                    self.swap(i, l)
                return

            else:
                # has two children
                right = self.ordered[r]

                # heap property maintained
                if min(left['score'], right['score']) >= current['score']:
                    return

                # left is smaller
                if left["score"] < right["score"]:
                    self.swap(i, l)
                    i = l

                else:
                    # right child is the smallest
                    self.swap(i, r)
                    i = r

    def swap(self, i, j):
        self.ordered[i], self.ordered[j] = self.ordered[j], self.ordered[i]
        map(self._update_map, [i, j])

    def _update_map(self, i):
        self.map[self.ordered[i]["key"]] = i

    def parent(self, i):
        return (i - 1) // 2

    def left_child(self, i):
        return 2 * i + 1

    def right_child(self, i):
        return 2 * i + 2

    def one_child(self, i):
        return self.right_child(i) >= self.length


def dijkstra(G, v):
    heap = BinaryHeap()
    heap.insert(v, 0)
    final_dist = {}

    while heap.length:

        # Found shortest path to node w.
        # lock it down!
        smallest_so_far = heap.pop()
        w = smallest_so_far['key']
        final_dist[w] = smallest_so_far['score']

        for x in G[w]:

            # Skip nodes we already have shortest
            # dist for.
            if x in final_dist:
                continue

            new_dist = final_dist[w] + G[w][x]

            if x not in heap.map:

                # add new node to the heap
                heap.insert(x, new_dist)

            elif new_dist < heap.fetch_score(x):

                # found a better path to the node
                heap.decrement_score(x, new_dist)

    return final_dist


############
#
# Test


def make_link(G, node1, node2, w):
    if node1 not in G:
        G[node1] = {}
    if node2 not in G[node1]:
        (G[node1])[node2] = 0
    (G[node1])[node2] += w
    if node2 not in G:
        G[node2] = {}
    if node1 not in G[node2]:
        (G[node2])[node1] = 0
    (G[node2])[node1] += w
    return G


def test():
    # shortcuts
    (a, b, c, d, e, f, g) = ('A', 'B', 'C', 'D', 'E', 'F', 'G')
    triples = ((a, c, 3), (c, b, 10), (a, b, 15), (d, b, 9), (a, d, 4), (d, f, 7), (d, e, 3),
               (e, g, 1), (e, f, 5), (f, g, 2), (b, f, 1))
    G = {}
    for (i, j, k) in triples:
        make_link(G, i, j, k)

    dist = dijkstra(G, a)
    assert dist[g] == 8  # (a -> d -> e -> g)
    assert dist[b] == 11  # (a -> d -> e -> g -> f -> b)

    # TODO: generate a random graph
    # for 100 nodes, pick a random node,
    # pick a random weight make a connection


if __name__ == '__main__':
    test()