CODEWITHSUNDEEP
pythonsortingoop
DBedrenko
DBedrenko

Reputation: 5029

How to create custom sort function with complicated logic?

I'm try to sort a list of objects based on some non-trivial comparison logic, but finding it difficult because in Python the custom sort function takes only 1 argument. In Java, for example, the sort function would have references to object1 and object2, making it straightforward to compare them.

class Point:
    def __init__(self, char, num, pt_type):
        self.char = char
        self.num = num
        self.pt_type = pt_type  # 'start' or 'end'

    def __str__(self):
        return str([self.char, str(self.num), self.pt_type])
    def __repr__(self):
        return str(self)

arr = [Point('C', 1, 'end'), Point('C', 9, 'start'),
       Point('B', 7, 'end'), Point('B', 2, 'end'),
       Point('A', 3, 'start'), Point('A', 6, 'start')]

def my_sort(key):
    # Sort by first element (letter). 
    #
    # If the letter is the same, fallback to sorting by the
    # 2nd element (number), but the logic of this comparison depends
    # on `pt_type`:
    #   -If Point1 and Point2 both have type 'start', pick the higher number first.
    #   -If Point1 and Point2 both have type 'end', pick the lower number first.
    #   -If Point1 and Point2 have different types, pick the 'start' type first.
    return key.char

print(sorted(arr, key=my_sort))

The expected sorted order should be:

[Point('A', 6, 'start'), Point('A', 3, 'start')
 Point('B', 2, 'end'), Point('B', 7, 'end'),
 Point('C', 9, 'start'), Point('C', 1, 'end')]

I don't know how to even start implementing the required logic, so I would be grateful for a push in the right direction.

Upvotes: 1

Views: 347

Answers (3)

jpp
jpp

Reputation: 164663

You can make sorting a property of your class, then use sorted. The benefit of this method: for no additional effort, you are able to compare objects with each other via comparison operators such as >, <, ==.

Specify __eq__ and __lt__ methods

At a minimum you should specify __eq__ and __lt__ methods:

class Point:
    def __init__(self, char, num, pt_type):
        self.char = char
        self.num = num
        self.pt_type = pt_type  # 'start' or 'end'

    def __str__(self):
        return str([self.char, str(self.num), self.pt_type])

    def __repr__(self):
        return str(self)

    def __eq__(self, other):
        return self.char == other.char and self.pt_type == other.pt_type

    def __lt__(self, other):
        if self.char != other.char:
            return self.char < other.char
        if (self.pt_type == 'start') and (other.pt_type == 'start'):
            return self.num > other.num
        elif (self.pt_type == 'end') and (other.pt_type == 'end'):
            return self.num < other.num
        else:
            return self.pt_type == 'start'

Adding other comparison methods such as __gt__, __ge__, etc, may be simplified via functools.total_ordering:

from functools import total_ordering

@total_ordering
class Point:
    def __init__(self, ...):
        # initialization logic
    def __eq__(self, other):
        # as before
    def __lt__(self, other):
        # as before

Example

arr = [Point('C', 1, 'end'), Point('C', 9, 'start'),
       Point('B', 7, 'end'), Point('B', 2, 'end'),
       Point('A', 3, 'start'), Point('A', 6, 'start')]

print(sorted(arr))

[['A', '6', 'start'],
 ['A', '3', 'start'],
 ['B', '2', 'end'],
 ['B', '7', 'end'],
 ['C', '9', 'start'],
 ['C', '1', 'end']]

Upvotes: 3

Dani Mesejo
Dani Mesejo

Reputation: 61910

I would use the following key function:

class Point:
    def __init__(self, char, num, pt_type):
        self.char = char
        self.num = num
        self.pt_type = pt_type  # 'start' or 'end'

    def __str__(self):
        return str([self.char, str(self.num), self.pt_type])

    def __repr__(self):
        return str(self)


arr = [Point('C', 1, 'end'), Point('C', 9, 'start'),
       Point('B', 7, 'end'), Point('B', 2, 'end'),
       Point('A', 3, 'start'), Point('A', 6, 'start')]


def key(p):
    return p.char, int(p.pt_type != 'start'), p.num if p.pt_type == 'end' else -1 * p.num


result = sorted(arr, key=key)
print(result)

Output

[['A', '6', 'start'], ['A', '3', 'start'], ['B', '2', 'end'], ['B', '7', 'end'], ['C', '9', 'start'], ['C', '1', 'end']]

The key function creates a tuple to be used as key, the first element is the letter, the second element is 0 if the node is of type 'start', 1 if is of type 'end'. The last element is negative if it is of type 'start', positive if it is of type 'end'.

Upvotes: 4

Geza Lore
Geza Lore

Reputation: 541

You want to use the cmp argument to sorted which takes a comparison function of 2 arguments: https://docs.python.org/2/library/functions.html#sorted

For your reference, the key function would compute a derived value from each item being sorted and sort according to that value, e.g. to sort a list of pairs by the second value in the pair you could do: sorted(items, key=lambda x: x[1])

Upvotes: 1

Related Questions