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