python - Calculate number of jumps in Dijkstra's algorithm? -

what fastest way in numpy calculate number of jumps dijkstra's algorithm uses? have 10000x10000 element connectivity matrix , use scipy.sparse.csgraph.dijkstra calculate filled distance matrix , predecessor matrix. naive solution follows:

import numpy np numpy.random import rand scipy.sparse.csgraph import dijkstra  def dijkway(dijkpredmat, i, j):     """calculate path between 2 nodes in dijkstra matrix"""     wayarr = []     while (i != j) & (j >= 0):         wayarr.append(j)         j = dijkpredmat[i,j]     return np.array(wayarr)  def jumpvec(pmat,node):     """calculate number of jumps 1 node others"""     jumps = np.zeros(len(pmat))     jumps[node] = -999     while 1:         try:             rvec = np.nonzero(jumps==0)[0]             r = rvec.min()             dway = dijkway(pmat, node, r)             jumps[dway] = np.arange(len(dway),0,-1)         except valueerror:             break     return jumps  #create matrix mat = (rand(500,500)*20) mat[(rand(50000)*500).astype(int), (rand(50000)*500).astype(int)] = np.nan dmat,pmat = dijkstra(mat,return_predecessors=true)  timeit jumpvec(pmat,300) in [25]: 10 loops, best of 3: 51.5 ms per loop 

~50msek/node ok expanding distance matrix 10000 nodes increases time ~2sek/node. jumpvec has executed 10000 times well...

following code can speedup 4x on pc, it's faster because:

• use ndarray.item() values array.
• use set object save unprocessed index.
• don't create numpy.arange() in while loop.

python code:

def dijkway2(dijkpredmat, i, j):     wayarr = []     while (i != j) & (j >= 0):         wayarr.append(j)         j = dijkpredmat.item(i,j)     return wayarr  def jumpvec2(pmat,node):     jumps = np.zeros(len(pmat))     jumps[node] = -999     todo = set()     in range(len(pmat)):         if != node:             todo.add(i)      indexs = np.arange(len(pmat), 0, -1)     while todo:         r = todo.pop()         dway = dijkway2(pmat, node, r)         jumps[dway] = indexs[-len(dway):]         todo -= set(dway)     return jumps 

to speedup more, can use cython:

import numpy np cimport numpy np import cython  @cython.wraparound(false) @cython.boundscheck(false) cpdef dijkway3(int[:, ::1] m, int i, int j):     cdef list wayarr = []     while (i != j) & (j >= 0):         wayarr.append(j)         j = m[i,j]     return wayarr  @cython.wraparound(false) @cython.boundscheck(false) def jumpvec3(int[:, ::1] pmat, int node):     cdef np.ndarray jumps     cdef int[::1] jumps_buf     cdef int i, j, r, n     cdef list dway     jumps = np.zeros(len(pmat), int)     jumps_buf = jumps     jumps[node] = -999      in range(len(jumps)):         if jumps_buf[i] != 0:             continue         r =         dway = dijkway3(pmat, node, r)         n = len(dway)         j in range(n):             jumps_buf[<int>dway[j]] = n - j     return jumps 

here test, cython version 80x faster:

%timeit jumpvec3(pmat,1) %timeit jumpvec2(pmat, 1) %timeit jumpvec(pmat, 1) 

output:

1000 loops, best of 3: 138 µs per loop 100 loops, best of 3: 2.81 ms per loop 100 loops, best of 3: 10.8 ms per loop