Improved stable marriage problem Python implementation

CONTRIBUTORS.md

@@ -51,3 +51,4 @@ This file lists everyone, who contributed to this repo and wanted to show up her
 - Vincent Zalzal
 - Jonathan D B Van Schenck
 - James Goytia
+- Amaras

contents/stable_marriage_problem/code/python/stable_marriage.py

@@ -1,25 +1,19 @@
 # Submitted by Marius Becker
+# Updated by Amaras
+
 
-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,62 +25,64 @@ 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:
+
+    while True:
         # Let every man without a partner propose to a woman
         for man in men:
-            if not man.has_partner():
+            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
+        if all((man.has_partner for man in men)):
+            return
+
 
 class Person:
-    name = None
-    preference = None
-    pref_index = 0
-    candidates = None
-    partner = None
 
     def __init__(self, name):
         self.name = name
         self.preference = []
         self.candidates = []
+        self.pref_index = 0
+        self._partner = None
 
-    def get_next_choice(self):
+    @property
+    def next_choice(self):
         """Return the next person in the own preference list"""
-        if self.pref_index >= len(self.preference):
+        try:
+            return self.preference[self.pref_index]
+        except IndexError:
             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 = self.next_choice
         person.candidates.append(self)
         self.pref_index += 1
 
@@ -99,33 +95,45 @@ def pick_preferred(self):
             if person == self.partner:
                 break
             elif person in self.candidates:
-                self.set_partner(person)
+                self.partner = person
                 break
 
-        # Rejected candidates don't get a second chance. :(
+        # Rejected candidates don't get a second chance
         self.candidates.clear()
 
-    def get_partner(self):
-        """Return the current partner"""
-        return self.partner
+    @property
+    def partner(self):
+        return self._partner
+
+    # The call self.partner = person sets self._partner as person
+    # However, since engagement is symmetrical, self._partner._partner
+    # (which is then person._partner) also needs to be set to self
+    @partner.setter
+    def partner(self, person):
+        """Set a person as the new partner and sets the partner of that
+        person as well"""
 
-    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:
+        if person != self._partner:
             # Remove self from current partner
-            if self.partner is not None:
-                self.partner.partner = None
+            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
+            self._partner = person
+            if self._partner is not None:
+                self._partner._partner = self
 
+    # This allows use of self.has_partner instead of self.has_partner()
+    @property
     def has_partner(self):
-        """Determine whether this person currently has a partner or not"""
-        return self.partner != None
+        """Determine whether this person currently has a partner or not."""
+        return self.partner is not 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()