I'm implementing the Dijkstra search algorithm in Python. At the end of the search, I reconstruct the shortest path using a predecessor map, starting with the destination node's predecessor. For example:
path = []
path.append(destination)
previous = predecessor_map[destination]
while previous != origin:
path.append(previous)
previous = predecessor_map[previous]
Is there any way to do this with less lines of code (e.g. list comprehension)?
The only suggestion that I have is to get rid of the slight code duplication:
path = []
previous = destination
while previous != origin:
path.append(previous)
previous = predecessor_map[previous]
Beyond that, I think your code is actually very clear and is unlikely to benefit from any attempts to shorten it.
Lastly, it is worth noting that the above also works when destination == origin, whereas your original version most probably doesn't (depends on how exactly predecessor_map is populated). Don't know if this is relevant to your use cases.
This might work:
path = [destination]
path += iter(lambda: predecessor_map[path[-1]], origin)
It behaves the same as your original code. But what you've already written is fine as is.
If destination could be equal to origin:
path = []
path += iter(lambda: predecessor_map[path[-1]] if path else destination, origin)
It behaves the same as @aix's code.
def backwalk(mymap, start, origin):
yield start
current = mymap[start]
while current != origin:
yield current
current = mymap[current]
path = list(backwalk(predecessor_map, destination, origin))
This separates the walking and collecting tasks.
If you can ensure that you never start with the origin, you can simplify to
def backwalk(mymap, start, origin):
current = start
while current != origin:
yield current
current = mymap[current]
You can recursively traverse the edges assuming predecessor_map is a dict mapping node to parent node and that None is the root:
predecessor_map={0: None, 1: None, 2: 1, 3: 1, 4: 0, 5: 1}
Define a recursive function that traverses the tree in reverse:
def path(node, predecessors):
return [None] if node is None else [node] + path(predecessors.get(node), predecessors)
Or, if you dare, a Y combinator:
Y=lambda f: (lambda x: f(lambda *args: x(x)(*args)))(lambda x: f(lambda *args: x(x)(*args)))
path=Y(lambda f: lambda node, p: [None] if node is None else [node] + f(p.get(node), p))
In use (using list comprehension):
>>> print [node for node in path(None, predecessor_map)]
[None]
>>> print [node for node in path(0, predecessor_map)]
[0, None]
>>> print [node for node in path(1, predecessor_map)]
[1, None]
>>> print [node for node in path(2, predecessor_map)]
[2, 1, None]
>>> print [node for node in path(3, predecessor_map)]
[3, 1, None]
>>> print [node for node in path(4, predecessor_map)]
[4, 0, None]
>>> print [node for node in path(5, predecessor_map)]
[5, 1, None]
One more possible solution is to use functional style programming with deferred output:
from itertools import tee, chain, imap, takewhile
predecessor_map = {2:1, 3:2}
destination = 3
origin = 1
def backwalk(predecessor_map, start, origin):
def deffered_output():
for i in output:
yield i
result, a = tee(deffered_output())
b = imap(predecessor_map.get,a)
output = takewhile(lambda x: x!=origin,chain([start],b))
return result
print(list(backwalk(predecessor_map,destination,origin)))
I personally wouldn't use this approach. But it's interesting for training though.
Explanation
The key element is deferred_output which postpones the calls to output.
Then we split output into 2 iterators using tee.
Then we apply predecessor_map.get to the second iterator called a and assign the new iterator to b.
Then we control the output with takewhile and stop when origin is reached.
I don't think you can do this iteration with a comprehension. Maybe you could simplify it a little, like this:
path, previous = [], destination
while True:
path.append(previous)
previous = predecessor_map[previous]
if previous == origin:
break
The above loop would look nicer with a do..while , but Python lacks it
