Untitled

import heapq

class CityMap:
    def __init__(self, roads, tracks, friends):
        # Calculate total number of unique locations from roads and tracks
        self.total_places = max(max(route[0], route[1]) for route in roads + tracks) + 1

        # Setup lists for connections
        self.paths_tracks = [[] for _ in range(self.total_places)]
        self.paths_roads = [[] for _ in range(self.total_places)]
        self.locations_friends = [None] * self.total_places

        # Fill in the track paths from the tracks data
        for start, end, weight in tracks:
            self.paths_tracks[start].append((end, weight))

        # Setup road connections, ensuring bi-directional nature
        for start, end, weight in roads:
            self.paths_roads[start].append((end, weight))
            self.paths_roads[end].append((start, weight))
        
        # Map friends to their respective locations
        for friend, position in friends:
            self.locations_friends[position] = friend

    def find_shortest_path(self, origin, path_type="road"):
        # Initialize distances and previous node trackers for all locations
        min_distances = [float('inf')] * self.total_places
        path_predecessors = [-1] * self.total_places
        min_distances[origin] = 0
        priority_queue = [(0, origin)]  # Pair of (distance, location index)

        while priority_queue:
            current_dist, current_loc = heapq.heappop(priority_queue)

            if current_dist > min_distances[current_loc]:
                continue  # Skip if a better path has already been found

            # Select the appropriate connections based on path type
            if path_type == "road":
                active_connections = self.paths_roads
            else:
                active_connections = self.paths_tracks

            # Update distances for all connected nodes
            for adjacent, weight in active_connections[current_loc]:
                new_dist = current_dist + weight
                if new_dist < min_distances[adjacent]:
                    min_distances[adjacent] = new_dist
                    path_predecessors[adjacent] = current_loc
                    heapq.heappush(priority_queue, (new_dist, adjacent))

        return min_distances, path_predecessors

    def trace_route(self, route_predecessors, target):
        # This function traces back the route from a given target using predecessors array
        travel_path = []
        step = target
        while step != -1:
            travel_path.append(step)
            step = route_predecessors[step]
        return travel_path[::-1]

    def search_nearby_stations(self, friend_pos):
        # This method identifies all nearby stations from a friend's position within two stops
        checked = [False] * self.total_places  # Changed from self.num_locations to self.total_places
        process_queue = [(friend_pos, 0)]  # holds tuples of (current position, number of stops)
        nearby_stations = []
        
        while process_queue:
            place, stops = process_queue.pop(0)
            if stops > 2:
                break  # limit search to two stops
            nearby_stations.append((place, stops))
            checked[place] = True
            
            if stops < 2:
                for adjacent, _ in self.paths_tracks[place]:  # Ensure you're using the correct paths
                    if not checked[adjacent]:
                        process_queue.append((adjacent, stops + 1))

        return nearby_stations

    def plan(self, start, destination):
        # Calculate the shortest paths from start to all destinations on roads and tracks
        road_distances, road_predecessors = self.find_shortest_path(start, "road")
        track_distances, _ = self.find_shortest_path(start, "track")
        
        optimal_route = None
        minimum_time = float('inf')
        optimal_friend = None
        optimal_station = None
        optimal_stops = float('inf')

        # Check each friend's location for possible pickup options
        for loc, friend in enumerate(self.locations_friends):  # Changed from self.friend_locations to self.locations_friends
            if friend is None:
                continue
            
            possible_pickups = self.search_nearby_stations(loc)
            
            for pickup_loc, stops in possible_pickups:
                # Determine the best path to destination from the pickup location
                distances_to_destination, prev_nodes_to_dest = self.find_shortest_path(pickup_loc, "road")
                track_distances_to_dest, _ = self.find_shortest_path(pickup_loc, "track")

                # Combine times from start to pickup and pickup to destination
                time_via_road = road_distances[pickup_loc] + distances_to_destination[destination]
                time_via_track = track_distances[pickup_loc] + track_distances_to_dest[destination]

                combined_time = min(time_via_road, time_via_track)

                # Select the best route based on time and number of stops needed to pick up the friend
                if combined_time < minimum_time or (combined_time == minimum_time and stops < optimal_stops):
                    minimum_time = combined_time
                    optimal_route = self.trace_route(road_predecessors, pickup_loc) + \
                                    self.trace_route(prev_nodes_to_dest, destination)[1:]  # Avoid duplicating nodes
                    optimal_friend = friend
                    optimal_station = pickup_loc
                    optimal_stops = stops

        return minimum_time, optimal_route, optimal_friend, optimal_station