A few months ago, I worked through Spotify’s first puzzle, a straightforward binary reversal problem. Round two is harder: the classic selection problem.
The problem: given an array of values (in this case, songs with a quality score), find the top k values.
Approach: max-heap with heapq Link to heading
The cleanest solution is a binary max-heap: push all elements in, then pop exactly k times. Total time complexity: O(n log k).
Python’s heapq module implements a min-heap. To use it as a max-heap, I defined a custom Song class with an inverted comparison: a song is “less than” another if it has a higher quality:
# 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 "less than" another
# if it has greater quality (inverted comparison).
return self.quality > other.quality
def __eq__(self, other):
return self.quality == other.quality and \
self.index == other.index
Note: Python 3 deprecates __cmp__, so the comparison protocol uses __lt__ and __eq__ instead.
Testing and submission Link to heading
As before, I used pytest to validate the solution against test cases before submission:
The automated Spotify judge accepted the solution:
Code is available on GitHub.