Triathlon Scheduling (solution from Kleinberg and Tardos' book on algorithm design)

triathlon.py

#!/usr/bin/env python3

from itertools import permutations
from random import randint

def tournamentTime(person_list):
     """ Computes how much a tournament will last according to the order of
     persons in 'person_list'
     Assumes a person is described as a 3-tuple of positive integers (ts, tb, tr) 
     where:
     ts: swimming time
     tc: cycling time
     tr: running time

     """
     if len(person_list) == 0:
          return 0
     startTime = 0
     finishTimes = []
     for p in person_list:
          startTime += p[0]
          finishTimes.append(startTime + p[1] + p[2])
     return max(finishTimes)

def bruteForce(person_set):
     """ Brute force algorithm to solve the Triathlon Scheduling problem.
     It first generates all the possible orderings and applies the function
     'tournamentTime' to compute the smallest possible tournament time.
     It receives a set and works for len(person_set) < 10 approximately.

     """

     if len(person_set) == 0:
          return 0

     perm = permutations(person_set)
     return min(map(tournamentTime, perm))

def greedy(person_set):
     """ Greedy algorithm to solve the Triathlon Scheduling problem.
     This algorithm sorts the persons acording to the sum of their cycling and 
     running times in reverse order.

     """
     if len(person_set) == 0:
          return 0

     person_list = list(person_set)
     person_list.sort(key=lambda x:x[1]+x[2], reverse=True)
     return tournamentTime(person_list)

def pgenerator(size, max_time = 100):
     """ Generates a set with 'size' persons and their timings.
     Each timing in limited by 'max_time'. If no max_time is given
     assume the default value 100.

     """

     if (size < 1 or max_time < 1):
          raise ValueError("Wrong size or max_time passed as parameter.")

     pset = set()
     while size > 0:
          lp = len(pset)
          p = (randint(1, max_time), randint(1, max_time), randint(1, max_time))
          pset.add(p)
          if len(pset) != lp:
               size -= 1
     return pset

def main():
     pset = pgenerator(100)
     print(greedy(pset))

if __name__ == "__main__":
    main()