Flight Ticket Booking Problem

solution.py

# -*- coding: utf-8 -*-
import datetime
from collections import defaultdict

FLIGHTS = {
    "chengdu_to_xian": {
        "reward": [
            {'name': 'GD2501', 'time': '08:00', 'price': 500, 'at_weekends': True},
            {'name': 'GD2501', 'time': '08:00', 'price': 800, 'at_weekends': False},
            {'name': 'GD2606', 'time': '12:25', 'price': 500, 'at_weekends': True},
            {'name': 'GD2606', 'time': '12:25', 'price': 1100, 'at_weekends': False},
            {'name': 'GD8732', 'time': '19:30', 'price': 400, 'at_weekends': True},
            {'name': 'GD8732', 'time': '19:30', 'price': 1000, 'at_weekends': False},
        ],
        "regular": [
            {'name': 'GD2501', 'time': '08:00', 'price': 900, 'at_weekends': True},
            {'name': 'GD2501', 'time': '08:00', 'price': 1100, 'at_weekends': False},
            {'name': 'GD2606', 'time': '12:25', 'price': 600, 'at_weekends': True},
            {'name': 'GD2606', 'time': '12:25', 'price': 1600, 'at_weekends': False},
            {'name': 'GD8732', 'time': '19:30', 'price': 1500, 'at_weekends': True},
            {'name': 'GD8732', 'time': '19:30', 'price': 2200, 'at_weekends': False},
        ],
    },
    "xian_to_chengdu": {
        "reward": [
            {'name': 'GD2502', 'time': '12:00', 'price': 800, 'at_weekends': True},
            {'name': 'GD2502', 'time': '12:00', 'price': 800, 'at_weekends': False},
            {'name': 'GD2607', 'time': '16:25', 'price': 500, 'at_weekends': True},
            {'name': 'GD2607', 'time': '16:25', 'price': 1100, 'at_weekends': False},
            {'name': 'GD8733', 'time': '23:30', 'price': 400, 'at_weekends': True},
            {'name': 'GD8733', 'time': '23:30', 'price': 1500, 'at_weekends': False},
        ],
        "regular": [
            {'name': 'GD2502', 'time': '12:00', 'price': 900, 'at_weekends': True},
            {'name': 'GD2502', 'time': '12:00', 'price': 1700, 'at_weekends': False},
            {'name': 'GD2607', 'time': '16:25', 'price': 600, 'at_weekends': True},
            {'name': 'GD2607', 'time': '16:25', 'price': 1600, 'at_weekends': False},
            {'name': 'GD8733', 'time': '23:30', 'price': 1000, 'at_weekends': True},
            {'name': 'GD8733', 'time': '23:30', 'price': 1600, 'at_weekends': False},
        ],
    }
}


class FlightFinder(object):
    def __init__(self, customer_type, departing_date, returning_date):
        self.customer_type = customer_type
        self.departing_date = departing_date
        self.returning_date = returning_date

    def run(self):
        departing_flights = FLIGHTS['chengdu_to_xian'][self.customer_type]
        returning_flights = FLIGHTS['xian_to_chengdu'][self.customer_type]
        return self.find_way(departing_flights, returning_flights)

    def find_way(self, departing_flights, returning_flights):
        min_amount = float('+inf')
        amount_map = defaultdict(list)

        for df in departing_flights:
            for rf in returning_flights:
                if self.is_available_way(df, rf):
                    amount = self.get_amount(df, rf)
                    if amount <= min_amount:
                        min_amount = amount
                        amount_map[min_amount].append((df, rf))

        flights = amount_map.get(min_amount)
        if flights:
            return flights[-1]  # FLIGHTS有序，返回最后一个

    def is_available_way(self, departing_flight, returning_flight):
        start_time = self.earliest_time(departing_flight)
        return_time = self.latest_time(returning_flight)
        if start_time and return_time and start_time < return_time:
            return True

    def earliest_time(self, flight):
        if flight['at_weekends']:
            date = DateHelper.next_weekend(self.departing_date)
        else:
            date = DateHelper.next_weekday(self.departing_date)
        if date <= self.returning_date:
            hour, minute = map(lambda x: int(x), flight['time'].split(':'))
            return datetime.datetime.combine(date, datetime.time(hour, minute))

    def latest_time(self, flight):
        if flight['at_weekends']:
            date = DateHelper.prev_weekend(self.returning_date)
        else:
            date = DateHelper.prev_weekday(self.returning_date)
        if date >= self.departing_date:
            hour, minute = map(lambda x: int(x), flight['time'].split(':'))
            return datetime.datetime.combine(date, datetime.time(hour, minute))

    @staticmethod
    def get_amount(departing_flight, returning_flight):
        return departing_flight['price'] + returning_flight['price']


class DateHelper(object):
    @classmethod
    def next_weekday(cls, date):
        weekday = date.weekday()
        if weekday < 5:
            return date
        else:
            return date + datetime.timedelta(days=7 - weekday)

    @classmethod
    def next_weekend(cls, date):
        weekday = date.weekday()
        if weekday > 4:
            return date
        else:
            return date + datetime.timedelta(days=5 - weekday)

    @classmethod
    def prev_weekday(cls, date):
        weekday = date.weekday()
        if weekday < 5:
            return date
        else:
            return date + datetime.timedelta(days=4 - weekday)

    @classmethod
    def prev_weekend(cls, date):
        weekday = date.weekday()
        if weekday > 4:
            return date
        else:
            return date + datetime.timedelta(days=-1 - weekday)


if __name__ == '__main__':
    flight_finder = FlightFinder('reward',
                                 datetime.date(2016, 04, 10),
                                 datetime.date(2016, 04, 20))
    flights = flight_finder.run()
    if flights:
        departing_flight, returning_flight = flights
        print departing_flight['name']
        print returning_flight['name']
    else:
        print 'cannot find flights'