Amazon Online Assessment (OA) - Favorite Genres

Given a map Map<String, List<String>> userSongs with user names as keys and a list of all the songs that the user has listened to as values.

Also given a map Map<String, List<String>> songGenres, with song genre as keys and a list of all the songs within that genre as values. A song can only belong to one genre.

The task is to return a map Map<String, List<String>>, where the key is a user name and the value is a list of the user's favorite genre(s). Favorite genre is the most listened to genre. A user can have more than one favorite genre if he/she has listened to the same number of songs per each of the genres.

Examples

Example 1:

Input:

userSongs = {
 "David": ["song1", "song2", "song3", "song4", "song8"],
 "Emma":  ["song5", "song6", "song7"]
}
songGenres = {
    "Rock":    ["song1", "song3"],
    "Dubstep": ["song7"],
    "Techno":  ["song2", "song4"],
    "Pop":     ["song5", "song6"],
    "Jazz":    ["song8", "song9"]
}

Output:

{
 "David": ["Rock", "Techno"],
 "Emma":  ["Pop"]
}

Explanation:

David has 2 Rock, 2 Techno and 1 Jazz songs. So, he has 2 favorite genres. Emma has 2 Pop and 1 Dubstep songs. Pop is Emma's favorite genre.

Example 2:

Input:

userSongs = {
 "David": ["song1", "song2"],
 "Emma":  ["song3", "song4"]
},
songGenres = {}

Output:

{
 "David": [],
 "Emma":  []
}

Try it yourself

from collections import Counter
from heapq import heapify, heappop
from typing import Dict, List


def top(items):
    h = [(-c, s) for s, c in Counter(items).items()]
    if len(h) == 0:
        return
    heapify(h)
    tc, ts = heappop(h)
    yield ts
    while h:
        c, s = heappop(h)
        if c != tc:
            break
        yield s


def favorite_genres(
        user_songs: Dict[str, List[str]],
        genre_songs: Dict[str, List[str]],
) -> Dict[str, List[str]]:
    '''
    >>> favorite_genres({
    ...     "David": ["song1", "song2", "song3", "song4", "song8"],
    ...     "Emma":  ["song5", "song6", "song7"]
    ... }, {
    ...     "Rock":    ["song1", "song3"],
    ...     "Dubstep": ["song7"],
    ...     "Techno":  ["song2", "song4"],
    ...     "Pop":     ["song5", "song6"],
    ...     "Jazz":    ["song8", "song9"]
    ... })
    {'David': ['Rock', 'Techno'], 'Emma': ['Pop']}
    >>> favorite_genres({
    ...     "David": ["song1", "song2"],
    ...     "Emma":  ["song3", "song4"]
    ... }, {})
    {'David': [], 'Emma': []}
    '''

    song_genres = {
        s: g
        for g, ss in genre_songs.items()
        for s in ss
    }
    return {
        u: list(top(song_genres[s] for s in ss))
        for u, ss in user_songs.items()
    }


if __name__ == '__main__':
    import doctest
    doctest.testmod()