diff --git a/contents/stable_marriage_problem/code/python/stable_marriage.py b/contents/stable_marriage_problem/code/python/stable_marriage.py index 8271e9b4e..b892e687b 100644 --- a/contents/stable_marriage_problem/code/python/stable_marriage.py +++ b/contents/stable_marriage_problem/code/python/stable_marriage.py @@ -1,25 +1,19 @@ # Submitted by Marius Becker +# Updated by Amaras +# Updated again by Jonathan Dönszelmann -import sys from random import shuffle from copy import copy -from string import ascii_uppercase +from string import ascii_uppercase, ascii_lowercase -def main(): - # Set this to however many men and women you want - if len(sys.argv) > 1: - num_pairs = int(sys.argv[1]) - else: - num_pairs = 5 - # There are only 26 possible names - if num_pairs > 13: - print('You can\' have more than 13 pairs.') - return +def main(): + # Set this to however many men and women you want, up to 26 + num_pairs = 5 # Create all Person objects - men = [ Person(name) for name in ascii_uppercase[0:num_pairs] ] - women = [ Person(name) for name in ascii_uppercase[num_pairs:num_pairs*2] ] + men = [Person(name) for name in ascii_uppercase[:num_pairs]] + women = [Person(name) for name in ascii_lowercase[:num_pairs]] # Set everyone's preferences for man in men: @@ -31,101 +25,58 @@ def main(): shuffle(woman.preference) # Run the algorithm - resolve(men, women) + stable_marriage(men, women) # Print preferences and the result + print('Preferences of the men:') for man in men: - print('{}: {}'.format(man.name, ', '.join([ p.name for p in man.preference ]))) + print(man) + + print() + print('Preferences of the women:') for woman in women: - print('{}: {}'.format(woman.name, ', '.join([ p.name for p in woman.preference ]))) + print(woman) - print('') + print('\n') + print('The algorithm gave this solution:') for man in men: - print('{} + {}'.format(man.name, man.partner.name)) + print(f'{man.name} + {man.partner.name}') -def resolve(men, women): + +def stable_marriage(men, women): """Finds pairs with stable marriages""" - cont = True - while cont: - # Let every man without a partner propose to a woman - for man in men: - if not man.has_partner(): - man.propose_to_next() - - # Let the women pick their favorites - for woman in women: - woman.pick_preferred() - - # Continue only when someone is still left without a partner - cont = False - for man in men: - if not man.has_partner(): - cont = True - break -class Person: - name = None - preference = None - pref_index = 0 - candidates = None - partner = None + people = [*men, *women] + + for person in people: + if person.partner is None: + # find someone in this person's preferences who + # doesn't already have a partner. + for possible_partner in person.preference: + if possible_partner.partner is None: + partner = possible_partner + break + else: + raise Exception(f"Couldn't find a partner for {person}") + else: + partner = person.partner + person.partner = partner + partner.partner = person + + +class Person: def __init__(self, name): self.name = name self.preference = [] - self.candidates = [] - - def get_next_choice(self): - """Return the next person in the own preference list""" - if self.pref_index >= len(self.preference): - return None - - return self.preference[self.pref_index] - - def propose_to_next(self): - """Propose to the next person in the own preference list""" - person = self.get_next_choice() - person.candidates.append(self) - self.pref_index += 1 - - def pick_preferred(self): - """Pick a new partner or stay with the old one if they are preferred""" - # Iterate own preferences in order - for person in self.preference: - # Pick the first person that's either a new candidate or the - # current partner - if person == self.partner: - break - elif person in self.candidates: - self.set_partner(person) - break - - # Rejected candidates don't get a second chance. :( - self.candidates.clear() - - def get_partner(self): - """Return the current partner""" - return self.partner - - def set_partner(self, person): - """Set a person as the new partner and run set_partner() on that person - as well""" - # Do nothing if nothing would change - if person != self.partner: - # Remove self from current partner - if self.partner is not None: - self.partner.partner = None - - # Set own and the other person's partner - self.partner = person - if self.partner is not None: - self.partner.partner = self - - def has_partner(self): - """Determine whether this person currently has a partner or not""" - return self.partner != None + self.partner = None + + # This allows the preferences to be printed more elegantly + def __str__(self): + return f'{self.name}: {", ".join(p.name for p in self.preference)}' + if __name__ == '__main__': main()