duty_generation.py

#!/usr/bin/env python

from collections import defaultdict
from pprint import pprint
import random


def generate_trips(route_names, start_hour, end_hour):
  results = []
  for route_name in route_names:
    for hour in range(start_hour, end_hour):
      results.append((route_name, hour))
  return results


def generate_duties(num_duties, trips):
  # Create some data structues for each lookup of:
  # - The possible hours of the day
  # - The possible routes for each hour
  possible_hours = {t[1] for t in trips}
  possible_routes_for_hour = defaultdict(list)
  for t in trips:
    possible_routes_for_hour[t[1]].append(t[0])

  duties = []

  for _ in range(num_duties):
    num_trips = random.randrange(1, len(possible_hours) + 1)
    hours = random.sample(possible_hours, num_trips)
    duty = []
    for hour in sorted(hours):
      route_name = random.choice(possible_routes_for_hour[hour])
      duty.append((route_name, hour))
    duties.append(duty)

  return duties


def main():
  routes = 'ABCD'
  start_hour = 9
  end_hour = start_hour + 12

  trips = generate_trips(routes, start_hour, end_hour)
  duties = generate_duties(5, trips)
  pprint(duties)


if __name__ == '__main__':
  main()