The optimized way to find the top k leading candidates from a unsorted hash map in python

Question

I need to write a method that will accept a timestamp and takes an array of votes, I need to return the k leading candidates with that timestamp. I came up with the following solution,

Here is the input to the methods:

votes = [{'candidate':'a', 'timestamp':2},{'candidate':'c', 'timestamp': 5},{'candidate':'c', 'timestamp': 12}]

timestamp = 5

k = 5

And the method to solve the problem,

def leading_candidates(votes, timestamp,k):

    candidates = {}
    leading_candidates = []

    for vote in votes:
        if vote['timestamp'] <= timestamp:
            if vote['candidate'] not in candidates:
                candidates[vote['candidate']] = 1
            else:
                candidates[vote['candidate']] += 1

    sorted_votes = sorted(candidates.values())[:k]

    for candidate in candidates:
        if candidates[candidate] in sorted_votes:
            leading_candidates.append(candidate)

    return leading_candidates    

print(leading_candidates(votes, timestamp, 2))

As you can see the second solution has a time complexity of \$O(k\,n)\$ where k is the time it takes to find the index in the leading candidates sorted array, In the worst case, it can be \$O(n^2)\$ and because of sorting it may be at least \$O(n\,\log n)\$.

Is there any way we can make it work with \$O(n)\$?

Answer 1

So, you want to count something and afterwards get the top k? That sounds like a job for collections.Counter!

from collections import Counter

def leading_candidates(votes, timestamp, k):
    vote_counts = Counter(vote['candidate']
                          for vote in votes
                          if vote['timestamp'] <= timestamp)
    return [candidate[0] for candidate in vote_counts.most_common(k)]


if __name__ == "__main__":
    print(leading_candidates(votes, timestamp, 2))

This way you don't need to special case a candidate not yet having received a vote (something you could have also done with a collections.defaultdict(int)). And it is \$\mathcal{O}(n)\$.

Also note that if k is large, the line if candidates[candidate] in sorted_votes will become slow, as it is a linear scan. At the same time, you can iterate over the keys and values of a dictionary at the same time with candidates.items(), so you don't need to do candidates[candidate].

Python has an official style-guide, PEP8, which recommends using spaces after commas, which you forgot to do before k in the function signature.

You should always guard your code with an if __name__ == "__main__": guard to allow importing from the script without running it.

Answer 2

It took some time but finally, I understood how we can solve this in O (n).

Since the vote counts are Integers we can use bucket sort or in the worst case, we can use radix sort to sort the leading candidates and the resulting time complexity will be O(n).