import math
from collections import deque
from ywalk.types import Connection, Mode
class Graph:
WEIGHTS = {
'least-hop': lambda c: 1,
'least-time': lambda c: c.time + c.mode.weight,
}
def __init__(self):
self.connections = {}
self.nodes = 0
self.edges = 0
self.predicates = []
self.set_weight()
def add_predicate(self, predicate):
self.predicates.append(predicate)
def set_weight(self, spec=None):
self.weight_method = Graph.WEIGHTS[spec or 'least-time']
def add_connection(self, conn: Connection):
for place in [conn.origin, conn.destination]:
if place not in self.connections:
self.nodes = self.nodes + 1
self.connections[place] = set()
if conn not in self.connections[conn.origin]:
self.edges = self.edges + 1
self.connections[conn.origin].add(conn)
def add_recall(self, destination):
for origin in self.get_places():
if origin == destination:
continue
self.add_connection(Connection(origin, destination, Mode.RECALL, 0))
def __contains__(self, place):
return place in self.connections.keys()
def get_connections_from(self, place):
for conn in self.connections[place]:
yield conn
def get_connections_to(self, place):
for origin in self.get_places():
yield from self.get_connections_between(origin, place)
def get_connections_between(self, origin, destination):
for conn in self.connections[origin]:
if conn.destination == destination:
yield conn
def get_places(self):
return self.connections.keys()
def populate(self, gen):
for conn in gen:
self.add_connection(conn)
def shortest_paths(self, origin, destination=None):
tentative = list(self.get_places())
weights = {place: math.inf for place in self.get_places()}
hops = {}
weights[origin] = 0
while tentative:
current = min(tentative, key=weights.get)
tentative.remove(current)
# Exit early if we were given a destination and have reached it
if destination and current == destination:
return weights, hops
for conn in self.connections[current]:
if conn.destination not in tentative:
continue
next_hop = conn.destination
# Make sure to filter out any unwanted connections by setting
# their weight to infinity
if all(f(conn) for f in self.predicates):
alt_weights = weights[current] + self.weight_method(conn)
else:
alt_weights = math.inf
if alt_weights < weights[next_hop]:
weights[next_hop] = alt_weights
hops[next_hop] = conn
return weights, hops
def find_path(self, *args):
stops = deque(args)
current_stop = stops.popleft()
journey = []
while stops:
next_stop = stops.popleft()
_, hops = self.shortest_paths(current_stop, next_stop)
if not hops or next_stop not in hops:
return None
path = deque()
last_hop = next_stop
while last_hop in hops:
conn = hops[last_hop]
path.appendleft(conn)
last_hop = conn.origin
if path is None:
return None
journey.extend(path)
current_stop = next_stop
return journey
