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
else:
# 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.