Untitled

import heapq

class CityMap:
    def __init__(self, roads, tracks, friends):
        # Initialize the maximum index based on roads and tracks for sizing our lists
        self.total_places = max(max(u, v) for u, v, _ in roads + tracks) + 1

        # Create lists to store road and track connections
        self.paths_roads = [[] for _ in range(self.total_places)]
        self.paths_tracks = [[] for _ in range(self.total_places)]
        
        # Assign roads to their connections (no/both directions)
        for start, end, weight in roads:
            self.paths_roads[start].append((end, weight))
            self.paths_roads[end].append((start, weight))

        # Tracks (directional))
        for start, end, weight in tracks:
            self.paths_tracks[start].append((end, weight))
        
        # Initialize storage for friends' locations, linking them to their location indices
        self.locations_friends = [None] * self.total_places
        for friend, location in friends:
            self.locations_friends[location] = friend

    def find_shortest_path(self, origin, path_type="road"): # Dijsktra's Algorithm
        # Set initial distances to infinity and establish a starting point
        distances = [float('inf')] * self.total_places
        predecessors = [-1] * self.total_places
        distances[origin] = 0
        priority_queue = [(0, origin)]  # Using a heap queue to manage the frontier

        while priority_queue:
            # Extract the location with the smallest known distance
            dist, loc = heapq.heappop(priority_queue)
            if dist > distances[loc]:
                continue  # Found a better path already

            # Determine which set of connections to use
            connections = self.paths_roads if path_type == "road" else self.paths_tracks

            # Explore each connection emanating from the current location
            for neighbor, weight in connections[loc]:
                updated_dist = dist + weight
                if updated_dist < distances[neighbor]:
                    # Update the path to this neighbor if the new distance is smaller
                    distances[neighbor] = updated_dist
                    predecessors[neighbor] = loc
                    heapq.heappush(priority_queue, (updated_dist, neighbor))

        return distances, predecessors

    def trace_route(self, predecessors, destination):
        # Construct the route by backtracking from the destination using the predecessors list
        route = []
        current = destination
        while current != -1:
            route.append(current)
            current = predecessors[current]
        route.reverse()  # Reverse to get the route from origin to destination
        return route

    def search_nearby_stations(self, friend_location):
        # Discover nearby train stations from a friend's location within two train stops
        visited = [False] * self.total_places
        queue = [(friend_location, 0)]  # Start with the friend's location at zero stops
        stations_within_two_stops = []
        
        while queue:
            location, stops = queue.pop(0)
            if stops > 2:  # Stop if the train journey exceeds two stops
                continue
            stations_within_two_stops.append((location, stops))
            visited[location] = True  # Mark the current station as visited
            
            if stops < 2:
                # Explore further connections from current station
                for neighbor, _ in self.paths_tracks[location]:
                    if not visited[neighbor]:
                        queue.append((neighbor, stops + 1))

        return stations_within_two_stops


    def plan(self, start, destination):
        # Establish the shortest paths from the starting point using roads and tracks
        road_distances, road_predecessors = self.find_shortest_path(start, "road")
        track_distances, _ = self.find_shortest_path(start, "track")
        
        best_route = None
        best_time = float('inf')
        best_friend = None
        best_station = None
        fewest_stops = float('inf')

        # Evaluate pickup options at each friend's location
        for index, friend in enumerate(self.locations_friends):
            if friend is None:
                continue
            
            potential_pickups = self.search_nearby_stations(index)
            
            for pickup_location, number_of_stops in potential_pickups:
                # Calculate comprehensive paths from pickup points to the destination
                distances_from_pickup_to_dest, paths_from_pickup_to_dest = self.find_shortest_path(pickup_location, "road")
                track_distances_to_dest, _ = self.find_shortest_path(pickup_location, "track")

                # Determine the minimum combined travel time via road or track
                total_time_via_road = road_distances[pickup_location] + distances_from_pickup_to_dest[destination]
                total_time_via_track = track_distances[pickup_location] + track_distances_to_dest[destination]

                minimal_travel_time = min(total_time_via_road, total_time_via_track)

                # Update the optimal solution if a better one is found
                if minimal_travel_time < best_time or (minimal_travel_time == best_time and number_of_stops < fewest_stops):
                    best_time = minimal_travel_time
                    best_route = self.trace_route(road_predecessors, pickup_location) + \
                                self.trace_route(paths_from_pickup_to_dest, destination)[1:]  # Connect without repeating the node
                    best_friend = friend
                    best_station = pickup_location
                    fewest_stops = number_of_stops

        return best_time, best_route, best_friend, best_station