Last Case Passed A1 2004
import heapq class CityMap: def __init__(self, roads, tracks, friends): self.location_map = {} # Maps locations to indices for quick access all_locations = set(u for road in roads for u in road[:2]) all_locations.update(u for track in tracks for u in track[:2]) self.location_list = list(all_locations) self.locations = len(self.location_list) for index, location in enumerate(self.location_list): self.location_map[location] = index self.road_adj = [[] for _ in range(self.locations)] self.track_adj = [[] for _ in range(self.locations)] for u, v, w in roads: u_index = self.location_map[u] v_index = self.location_map[v] self.road_adj[u_index].append((v_index, w)) self.road_adj[v_index].append((u_index, w)) for u, v, w in tracks: u_index = self.location_map[u] v_index = self.location_map[v] self.track_adj[u_index].append((v_index, w)) self.friends = [(name, self.location_map[loc]) for name, loc in friends] def dijkstra(self, adj_list, start): dist = [float('inf')] * self.locations pred = [-1] * self.locations dist[start] = 0 pq = [(0, start)] visited = set() # To track visited nodes while pq: current_dist, u = heapq.heappop(pq) if u in visited: continue visited.add(u) for v, weight in adj_list[u]: if dist[v] > current_dist + weight: dist[v] = current_dist + weight pred[v] = u heapq.heappush(pq, (dist[v], v)) return dist, pred def plan(self, start, destination): if start == destination: for name, loc in self.friends: if loc == self.location_map[start]: return (0, [start], name, start) start_index = self.location_map[start] destination_index = self.location_map[destination] road_distances, road_pred = self.dijkstra(self.road_adj, start_index) best_time = float('inf') best_route = [] best_friend = None best_pickup_index = None for name, friend_loc in self.friends: track_distances, track_pred = self.dijkstra(self.track_adj, friend_loc) for i in range(self.locations): total_time = road_distances[i] + track_distances[i] if total_time < best_time: best_time = total_time route_from_start = self.reconstruct_path(road_pred, start_index, i) route_from_friend = self.reconstruct_path(track_pred, friend_loc, i) if route_from_friend and route_from_start[-1] == route_from_friend[0]: route_from_friend.pop(0) best_route = route_from_start + route_from_friend best_friend = name best_pickup_index = i # Final route to destination if best_route[-1] != destination: destination_route = self.reconstruct_path(road_pred, best_pickup_index, destination_index) if best_route[-1] == destination_route[0]: destination_route.pop(0) best_route += destination_route return best_time, best_route, best_friend, self.location_list[best_pickup_index] def reconstruct_path(self, pred, start, end): path = [] while end != start: path.append(self.location_list[end]) end = pred[end] path.append(self.location_list[start]) path.reverse() return path
Leave a Comment