Flight combinations between two cities

Question

I am trying to obtain all possible ways to go from a city A to a city B on a given day by plane, with a maximum of 2 stops.

I have as input a list of dictionaries:

flights = [
        ...
        {'dep': 'FRA', 'arr': 'AMS', 'dep_dt': '2017-05-01 12:00:00', 'arr_dt': '2017-05-01 13:15:00', 'price': 100},
        {'dep': 'FRA', 'arr': 'CPH', 'dep_dt': '2017-05-01 10:00:00', 'arr_dt': '2017-05-01 12:00:00', 'price': 80},
        {'dep': 'FRA', 'arr': 'MAD', 'dep_dt': '2017-05-01 09:00:00', 'arr_dt': '2017-05-01 10:50:00', 'price': 30},
        {'dep': 'CPH', 'arr': 'AMS', 'dep_dt': '2017-05-01 15:00:00', 'arr_dt': '2017-05-01 16:30:00', 'price': 60},
        {'dep': 'CPH', 'arr': 'MAD', 'dep_dt': '2017-05-01 14:15:00', 'arr_dt': '2017-05-01 17:10:00', 'price': 70},
        {'dep': 'MAD', 'arr': 'AMS', 'dep_dt': '2017-05-01 19:00:00', 'arr_dt': '2017-05-01 21:40:00', 'price': 20},
       ...
    ]

With thousands of records between many cities on many dates.

Example

Let's say I want all possibilities to go from FRA to AMS. To find all possible trips, I've split the problem in 3:

• Direct flights
• Flights with 1 stopover
• Flights with 2 stopovers

(I omit the date constraints below for simplicity)

In order to get the direct flights, I can do:

direct_flights = [f for f in flights if f['dep']=='FRA' and f['arr']=='AMS']

To get trips with one stopover, I've gone for a nested loop. Not optimal, but it works:

first_segment = [f for f in flights if f['dep']=='FRA']
second_segment = [f for f in flights if f['arr']=='AMS']

flights_w_1stop = []
for s1 in first_segment:
    for s2 in second_segment:
        if (s1['arr'] == s2['dep'] and
            s1['arr_dt'] < s2['dep_dt']):
            flights_w_1stop.append((s1, s2))

Now, to get all trips with two stopovers, I only could come up with a similar approach than above, but this makes two nested loops, which looks horrible and will be very slow when I have many flights.

first_segment = [f for f in flights if f['dep']=='FRA']
second_segment = [f for f in flights]
third_segment = [f for f in flights if f['arr']=='AMS']

flights_w_2stops = []
for s1 in first_segment:
    for s2 in second_segment:
        for s3 in third_segment:
            if (s1['arr'] == s2['dep'] and
                s2['arr'] == s3['dep'] and
                s1['arr_dt'] < s2['dep_dt'] and
                s2['arr_dt'] < s3['dep_dt']):
                flights_w_2stops.append((s1, s2, s3))

Finally, I would merge all 3 different lists with the same data structure and sort them either by price or by total duration.

Question

How could this be done more efficiently?

Note: I have read about graphs, but I don't think I will ever want to consider more than two stopovers. In addition, it would mean learning a new methodology. It is always good to learn, so if it is the best option I would invest the time. But if a similar thing can be done in another way (rearranging the nested loops somehow?), I would prefer so.

I found a similar question in SO and CodeReview, but I do not see how to apply that to my case.

Answer 1

Your code's runtime is \$O(n^3)\$, this is almost definitely not the performance that you want. The simplest way to solve your problem is by, you guessed it, a graph. Take the following directed graph of your selected flights:

Using a graph gives the ability to check only the actual routes. And so it wouldn't check if it is possible to fly from say; FRA -> CPH, and then FRA -> MAD, as the original code does. This benefit is good, however there is a larger benefit when there are say two, or more, flights in each edge. When that is the case, the original would go through roughly \$432\$ different arrangement of flights. Where a graph would only check the four routes to AMS, which would include travelling all \$6\$ edges. The graph could then return the \$1\$, or \$4\$, routes to AMS, and then go on to return all \$8\$, or \$18\$, permutations of those flights. And so the graph would check roughly \$28\$ permutations of these flights, as opposed to \$432\$.

Another example of this is if you wanted to obtain all the routes from from FRA -> MAD. And so you'd return the flights from FRA -> MAD and FRA -> CPH -> MAD. The first algorithm would still visit all \$6\$ airports, as you want to check if there are any three flights that can get to the desired destination. This time the code would return the above \$2\$ routes, which would lead to the graph quickly returning all \$6\$ possible flight combinations. Leading to roughly \$12\$ permutations.

Note: My maths is vague, and rough, it's not scientific, and depending on where it's measure, I'd not be suppressed if you obtained more permutations. It is merely to show the difference in magnitude between the two approaches.

There are two algorithms that I describe above, the first, can be implemented by:

1. Make an 'array' of the starting airport.

2. while 'the array' is not empty repeat the following commands:

   1. Pop a route from 'the array'.
   2. Get the last airports name from the route
   3. If the name is the destination yield it.
   4. If the route's length is equal to the depth we limit to, continue.
   5. Extend 'the array' with the last airports possible flights to.

And so could be:

def _routes(self, destination, depth=1):
    depth += 1
    l = [[self.name]]
    while l:
        route = l.pop()
        name = route[-1]
        if name == destination:
            yield route
        if len(route) == depth:
            continue
        l.extend([
            (route + [a])
            for a in self._graph[name]._flights.keys()
        ])

The second algorithm, is much simpler. You want to produce the product of the flights between each airport, for each route. Which can be:

def routes(self, destination, depth=1):
    root = self._graph
    return chain.from_iterable(
        product(*(
            root[dep].flights(arr)
            for dep, arr in pairwise(route)
        ))
        for route in self._routes(destination, depth)
    )

And so in total I'd use the below code. I added some un-used functions; remove_flight, and flights_to so that if you did want to use it rather than a list, you have a simpler interface.

from collections import defaultdict, namedtuple
from itertools import chain, product, tee

def pairwise(iterable):
    "s -> (s0,s1), (s1,s2), (s2, s3), ..."
    a, b = tee(iterable)
    next(b, None)
    return zip(a, b)

_DEFAULT = object()
class FlightGraph:
    Flight = namedtuple('Flight', 'dep arr dep_dt arr_dt price')
    class Airport:
        __slots__ = ['name', '_graph', '_flights']
        def __init__(self, graph, name):
            self._graph = graph
            self.name = name
            self._flights = defaultdict(set)

        def __getitem__(self, airport):
            if airport not in self._flights:
                raise KeyError('Airport has no flights to {!r}'.format(airport))
            return self._graph[airport]

        def __iter__(self):
            return iter(self._flights.items())

        def add_flight(self, airport, flight):
            # ensure airport exists
            self._graph[airport]
            self._flights[airport].add(flight)

        def remove_flight(self, airport, flight):
            self._flights[airport].remove(flight)

        def flights(self, airport, default=_DEFAULT):
            if airport in self._flights:
                return frozenset(self._flights[airport])
            if default is _DEFAULT:
                raise KeyError('Airport has no flights to {!r}'.format(airport))
            else:
                return default

        def flights_to(self):
            return tuple(self._flights.keys())

        def _routes(self, destination, depth=1):
            depth += 1
            l = [[self.name]]
            while l:
                route = l.pop()
                name = route[-1]
                if name == destination:
                    yield route
                if len(route) == depth:
                    continue
                l.extend([
                    (route + [a])
                    for a in self._graph[name]._flights.keys()
                ])

        def routes(self, destination, depth=1):
            root = self._graph
            return chain.from_iterable(
                product(*(
                    root[dep].flights(arr)
                    for dep, arr in pairwise(route)
                ))
                for route in self._routes(destination, depth)
            )

    def __init__(self, flights=[]):
        self._airports = {}
        for flight in flights:
            flight = self.Flight(**flight)
            self[flight.dep].add_flight(flight.arr, flight)

    def __getitem__(self, key):
        try:
            return self._airports[key]
        except KeyError:
            ret = self._airports[key] = self.Airport(self, key)
            return ret


flights = FlightGraph([
    {'dep': 'FRA', 'arr': 'AMS', 'dep_dt': '2017-05-01 12:00:00', 'arr_dt': '2017-05-01 13:15:00', 'price': 100},
    {'dep': 'FRA', 'arr': 'AMS', 'dep_dt': '2017-05-01 12:00:00', 'arr_dt': '2017-05-01 13:15:00', 'price': 150},
    {'dep': 'FRA', 'arr': 'CPH', 'dep_dt': '2017-05-01 10:00:00', 'arr_dt': '2017-05-01 12:00:00', 'price': 80},
    {'dep': 'FRA', 'arr': 'MAD', 'dep_dt': '2017-05-01 09:00:00', 'arr_dt': '2017-05-01 10:50:00', 'price': 30},
    {'dep': 'CPH', 'arr': 'AMS', 'dep_dt': '2017-05-01 15:00:00', 'arr_dt': '2017-05-01 16:30:00', 'price': 60},
    {'dep': 'CPH', 'arr': 'MAD', 'dep_dt': '2017-05-01 14:15:00', 'arr_dt': '2017-05-01 17:10:00', 'price': 70},
    {'dep': 'MAD', 'arr': 'AMS', 'dep_dt': '2017-05-01 19:00:00', 'arr_dt': '2017-05-01 21:40:00', 'price': 20},
])

for route in flights['FRA'].routes('AMS', 3):
    # TODO: filter if the arr_dt and dep_dt don't allow the route to be possible
    print(', '.join('{0.dep} -> {0.arr} ({0.price})'.format(f) for f in route))