Spotify puzzles: round two

Some months ago, I began challenging myself with Spotify puzzles: at that time I was dealing with an easy problem; now, the difficulty has increased. The round two consists in the typical “Selection problem“: given an array of values, find the max (or min) k values. I decided to still use Python and to use its heapq module to store values in a binary max-heap data structure, and then remove exactly k values from the top of the heap. This approach will guarantee that the total time complexity of the algorithm will be O(n log k).

The heapq module implements a min-heap data structure (as happens in Java). In order to use a max-heap, I specified a custom comparator for the Song class I wrote (remember: Python3 deprecates __cmp__ method, so I resort to implement __lt__ and __eq__ methods to specify a custom comparison algorithm):

# If two songs have the same quality, give precedence to the one
# appearing first on the album (presumably there was a reason for the
# producers to put that song before the other).
def __lt__(self, other):
    if self.quality == other.quality:
        return self.index < other.index
        # heapq is a min-heap, so a song is lt than another if it has
        # a greater quality (inverted!)
        return self.quality > other.quality

def __eq__(self, other):
    return self.quality == other.quality and \
           self.index == other.index

As always, I used py.test to automatically test my solution against test-cases.


In the end, the honorable (automated) judge confirms the solution:


Code has been published on GitHub.

Leave a Reply