Fast dijkstra

dijkstra.py

import csv
import heapq
from random import randint

# no heapq implementation
def up_heapify(length, P, L, i):
    idx, ix = L[i]
    while i:
        pi = (i-1) >> 1 #(i-1)/2
        pidx, pix = L[pi]
        if pidx > idx:
            P[ix], P[pix] = pi, i
            L[i], L[pi] = L[pi], L[i]
            i = pi
        else:
            return

def down_heapify(length, P, L, i):
    l = length[0]
    idx, ix = L[i]
    while True:        
        ri = (i+1) << 1 #2*i+2
        if ri >= l:
            li = ri-1
            if li >= l: return
            elif idx > L[li][0]: 
                P[ix], P[L[li][1]] = li, i
                L[i], L[li] = L[li], L[i]
            return
        else:
            li = ri-1
            lidx, lix = L[li]
            ridx, rix = L[ri]
            if lidx >= idx and ridx >= idx: return
            elif lidx < ridx:
                P[ix], P[lix] = li, i
                L[i], L[li] = L[li], L[i]
                i = li
            else:
                P[ix], P[rix] = ri, i
                L[i], L[ri] = L[ri], L[i]
                i = ri

def extract_min(length, P, L):
    li = length[0]-1
    P[L[li][1]] = 0
    minval = L[0]
    L[0] = L[li]
    length[0] -= 1
    down_heapify(length, P, L, 0)
    return minval

def heap_add(length, P, L, dx, x):
    li = length[0]
    L[li] = [dx, x]
    P[x] = li
    length[0] += 1
    up_heapify(length, P, L, li)

def dijkstra(G, v):
    dist_len, dist_pos, dist_heap = [0], {}, [None]*len(G)
    heap_add(dist_len, dist_pos, dist_heap, 0, v)
    final_dist = {}
    while dist_len[0]:
        dw, w = extract_min(dist_len, dist_pos, dist_heap)
        final_dist[w] = dw
        for x in G[w]:
            if x not in final_dist:
                dx = dw + G[w][x]
                if x not in dist_pos:
                    heap_add(dist_len, dist_pos, dist_heap, dx, x)
                else:
                    i = dist_pos[x]
                    if dist_heap[i][0] > dx:
                        dist_heap[i][0] = dx
                        up_heapify(dist_len, dist_pos, dist_heap, i)
    return final_dist

#-------------------------------------------
# Heap implementation
#
def H_heapify(H,H_pos,i):
  'Heapify the heap starting from node i'
  if i<1 or i>H[0]:
    raise ValueError
  start=i
  sdist,snode=H[i]
  while True:
    # Leaf node, no children
    if i>H[0]//2:        
      if i!=start:
        H[i]=(sdist,snode)
        H_pos[snode]=i
      return
    L_child=i*2
    R_child=i*2+1
    # Node with single left child
    if L_child==H[0]:    
      if H[L_child][0]<sdist:
        H[i]=H[L_child]
        H_pos[H[L_child][1]]=i
        i=L_child
      if i!=start:
        H[i]=(sdist,snode)
        H_pos[snode]=i
      return
    # Node with two children
    if H[L_child][0]<sdist:
      if H[R_child][0]<H[L_child][0]:
        H[i]=H[R_child]                 
        H_pos[H[R_child][1]]=i
        i=R_child
      else:
        H[i]=H[L_child]
        H_pos[H[L_child][1]]=i
        i=L_child
    elif H[R_child][0]<sdist:
      if H[L_child][0]<H[R_child][0]:
        H[i]=H[L_child]
        H_pos[H[L_child][1]]=i
        i=L_child
      else:
        H[i]=H[R_child]                 
        H_pos[H[R_child][1]]=i
        i=R_child
    else:
      if i!=start:
        H[i]=(sdist,snode)
        H_pos[snode]=i
      return

def H_heapify_UP(H,H_pos,i):
  'Heapify the heap upwards starting from node i'
  sdist,snode=H[i]
  start=i
  Parent=i//2
  while Parent and sdist<H[Parent][0]:
    H[i]=H[Parent]
    H_pos[H[Parent][1]]=i
    i=Parent
    Parent=i//2
  if i!=start:
    H[i]=(sdist,snode)
    H_pos[H[i][1]]=i

def H_build(H):
  'Builds heap from a list of tuples (node, distance) in place'
  H.insert(0,len(H))
  H_pos={}
  for i in range(1,H[0]+1):
    H_pos[H[i][1]]=i
  for i in range(H[0]//2,0,-1):
    H_heapify(H,H_pos,i)
  return H_pos

def H_remove_min(H,H_pos):
  'Removes the root node and returns its value'
  root=H[1][:]
  del H_pos[root[1]]
  H[1]=H[H[0]]
  H.pop()
  H[0]-=1
  if H[0]:
    H_pos[H[1][1]]=1
    H_heapify(H,H_pos,1)
  return root

def H_insert(H,H_pos,v):
  'Inserts tuple v into the heap'
  H.append(v)
  H[0]+=1
  i=H[0]
  Parent=i//2
  while Parent and v[0]<H[Parent][0]:
    H[i]=H[Parent]
    H_pos[H[Parent][1]]=i
    i=Parent
    Parent=i//2
  H[i]=v
  H_pos[H[i][1]]=i
#-------------------------------------------

def dijkstraMC(G,v):
  H=[(0,v)]
  H_pos=H_build(H)  # H[0] contains heap size
  distances={v:0}
  while H[0]:
    val,w=H_remove_min(H,H_pos)
    distances[w]=val
    for x in G[w]:
      if x not in distances:
        if x not in H_pos:
          H_insert(H,H_pos,(val+G[w][x],x))
        elif val+G[w][x]<H[H_pos[x]][0]:
          H[H_pos[x]]=(val+G[w][x],x)
          H_heapify_UP(H,H_pos,H_pos[x])
  return distances

############
### PQ implementation
############
REMOVED = -1                    # placeholder for a removed element

def add_item(pq, entry_finder, key, value):
    if key in entry_finder:
        remove_item(pq, entry_finder, key)
    entry = [value, key]
    entry_finder[key] = entry
    heapq.heappush(pq, entry)

def remove_item(pq, entry_finder, key):
    entry = entry_finder.pop(key)
    entry[-1] = REMOVED
    return entry[0], key

def pop_item(pq, entry_finder):
    while pq:
        value, key = heapq.heappop(pq)
        if key is not REMOVED:
            del entry_finder[key]
            return value, key

def dijkstra_heapq(G,v):
    pq = []                         # list of entries arranged in a heap
    entry_finder = {}               # mapping of tasks to entries
    add_item(pq, entry_finder, v, 0)
    final_dist = {}
    while len(entry_finder) and len(final_dist) < len(G):
        w_dist, w = pop_item(pq, entry_finder)
        # lock it down!
        final_dist[w] = w_dist
        for x in G[w]:
            if x not in final_dist:
                if x not in entry_finder:
                    add_item(pq, entry_finder, x, final_dist[w] + G[w][x])
                elif final_dist[w] + G[w][x] < entry_finder[x][0]:
                    add_item(pq, entry_finder, x, final_dist[w] + G[w][x])
    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 make_random_graph(size):
    G = {}
    nodes = range(size)
    for a in range(2*size):
        a = a % size
        b = (a + randint(1, size - 1)) % size
        w = randint(0, 20)
        make_link(G, a, b, w)
    while True:
        conn, unconn = connected_nodes(G, 0)
        if len(unconn) == 0:
            break
        make_link(G, conn.pop(), unconn.pop(), randint(0, 20))
    return G

def connected_nodes(G, root):
    # just run a search
    all_nodes = set(G.keys())
    open_list = [root]
    visited_nodes = set([root])
    all_nodes.remove(root)
    while len(open_list) > 0:
        node = open_list.pop()
        neighbors = G[node].keys()
        for n in neighbors:
            if n not in visited_nodes:
                open_list.append(n)
                visited_nodes.add(n)
                all_nodes.remove(n)
    return visited_nodes, all_nodes

def make_mc_graph(N):
    nodes=list(range(1,N))
    edges=[]
    for i in range(5*N):  # m=5*n at most
      u=randint(1,N)
      v=randint(1,N)
      if u!=v:
        edges.append((u,v,randint(1,10)))
    G = {}
    for (i,j,k) in edges:
        make_link(G, i, j, k)
    return G

def make_cliq_graph(N):
    G = {}
    for i in range(N):
        for j in range(i+1, N):
            make_link(G, i, j, randint(1, 100))
    return G

def test():
    import random
    # 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)
    print 'test passes'

def test2():
    (a,b,c,d) = ('A', 'B', 'C', 'D')
    triples = ((a,b,1),(a,c,4),(a,d,4),(b,c,1),(c,d,1))
    G = {}
    for (i,j,k) in triples:
        make_link(G, i, j, k)

    dist = dijkstra(G, a)

    assert dist[d] == 3 #(a -> b -> c -> d)
    print 'test2 passes'

def test_speed(make_my_graph):
    import time
    sizes = [100,300,600,1000,3000,6000,10000,30000]
    trials = 10
    print
    print 'using %s to generate random graphs' % (make_my_graph.__name__)
    for size in sizes:
        avg_time = 0.0
        avg_timeMC = 0.0
        avg_timeHQ = 0.0
        for _ in xrange(trials):
            G = make_my_graph(size)
            start_time = time.time()            
            dist = dijkstra(G, 1)
            avg_time += (time.time() - start_time)
            start_time = time.time()
            dist = dijkstraMC(G, 1)
            avg_timeMC += (time.time() - start_time)
            start_time = time.time()
            dist = dijkstra_heapq(G, 1)
            avg_timeHQ += (time.time() - start_time)
        avg_time /= trials
        avg_timeMC /= trials
        avg_timeHQ /= trials
        print 'avg of time taken for %d trials of %d size graph, avg time: %.5f, without heapq/without heapq MC: %.3f, without heapq/with heapq: %.3f' % (trials, size, avg_time, avg_time/avg_timeMC, avg_time/avg_timeHQ)

test()
test2()
test_speed(make_mc_graph)
test_speed(make_random_graph)
#test_speed(make_cliq_graph)