Count words in Python

count_words.py

"""Count words."""

def count_words(s, n):
    """Return the n most frequently occuring words in s."""
    
    # TODO: Count the number of occurences of each word in s
    words = s.split()
    word_occurences = [words.count(i) for i in words]
    word_map = dict(zip(words, word_occurences))
    # TODO: Sort the occurences in descending order (alphabetically in case of ties)
    # TODO: Return the top n words as a list of tuples (<word>, <count>)
    top_n =sorted(word_map.iteritems(),key=lambda(k,v):(-v,k))[0:n]
    return top_n


def test_run():
    """Test count_words() with some inputs."""
    print count_words("cat bat mat cat bat cat", 3)
    print count_words("betty bought a bit of butter but the butter was bitter", 3)


if __name__ == '__main__':
    test_run()