Search code examples
pythonargmax

Python - argmin / argmax for member function


I would like to know if there is a "pythonic" way to use the mathematical argmin/argmax for a member function without using library like numpy.

I have a class with a member function inside which returns an integer. I instantiate several objects of this class. I'd like to know which object has the lower return value for this method.

Please find below my source code. The part I'd like to improve is just after the tag Code I'd like to improve. This code is working very well, but I'm pretty sure there is a better way to do the same thing.

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""argmin example"""


class People(object):
    """People class"""
    ret_ages = {"Half": 60, "AlmostFull": 65, "Full": 71}  # years

    def __init__(self, name, age, ret_mode):
        super(People, self).__init__()

        if ret_mode not in self.ret_ages.keys():
            raise KeyError(ret_mode + " not in " + str(self.ret_ages.keys()))

        self.name = name
        self.age = age
        self.ret_mode = ret_mode

    def get_remaining_years(self):
        """
                Return how many years People have still to work before earning
                <rate> retirement.
                <rate> could be "Half", "Middle" or "Full".
        """
        try:
            return self.ret_ages[self.ret_mode] - self.age
        except KeyError:
            raise KeyError("rate has to be in " + str(self.ret_ages.keys()))


def main():
    """Main function"""
    people_list = [
        People("Juliette", 35, "Full"),
        People("Coralie", 26, "Half"),
        People("Laura", 27, "AlmostFull")
    ]

    # Debugging print
    for people in people_list:
        print people.name, "has still to work",\
            people.get_remaining_years(), "years."
    print
    # End of debugging print

    ############################
    # Code I'd like to improve #
    ############################

    people_closer_to_ret = people_list[0]
    minimum_remainining_years = people_closer_to_ret.get_remaining_years()

    for people in people_list:
        if people.get_remaining_years() < minimum_remainining_years:
            people_closer_to_ret = people
            minimum_remainining_years = people.get_remaining_years()

            minimum_remainining_years = people.get_remaining_years()

    ###################################
    # End of code I'd like to improve #
    ###################################

    print people_closer_to_ret.name, "will be retired soon !"


if __name__ == '__main__':
    main()

Here is the output of this script :

Juliette has still to work 36 years.
Coralie has still to work 34 years.
Laura has still to work 38 years.

Coralie will be retired soon !

Solution

  • A good way to write this code in a more pythonic way is to use min function, which can actually be used like argmin function thanks to its parameter key.

    If we replace the code between the tag Code I'd like to improve and End of code I'd like to improve by :

    people_closer_to_ret = min(people_list,
                               key=lambda people: people.get_remaining_years()),
    

    it's working perfectly. The key argument is useful for telling the min function which criterion it has to minimize.

    So your full code is the following :

    #!/usr/bin/env python
    # -*- coding: utf-8 -*-
    """argmin example"""
    
    
    class People(object):
        """People class"""
        ret_ages = {"Half": 60, "AlmostFull": 65, "Full": 71}  # years
    
        def __init__(self, name, age, ret_mode):
            super(People, self).__init__()
    
            if ret_mode not in self.ret_ages.keys():
                raise KeyError(ret_mode + " not in " + str(self.ret_ages.keys()))
    
            self.name = name
            self.age = age
            self.ret_mode = ret_mode
    
        def get_remaining_years(self):
            """
                    Return how many years People have still to work before earning
                    <rate> retirement.
                    <rate> could be "Half", "Middle" or "Full".
            """
            try:
                return self.ret_ages[self.ret_mode] - self.age
            except KeyError:
                raise KeyError("rate has to be in " + str(self.ret_ages.keys()))
    
    
    def main():
        """Main function"""
        people_list = [
            People("Juliette", 35, "Full"),
            People("Coralie", 26, "Half"),
            People("Laura", 27, "AlmostFull")
        ]
    
        # Debugging print
        for people in people_list:
            print people.name, "has still to work",\
                people.get_remaining_years(), "years."
        print
        # End of debugging print
    
        people_closer_to_ret = min(people_list,
                                   key=lambda people: people.get_remaining_years())
    
        print people_closer_to_ret.name, "will be retired soon !"
    
    
    if __name__ == '__main__':
        main()