[알고리즘] 프로그래머스 - 베스트앨범

프로그래머스 - 베스트앨범

내 풀이

from collections import defaultdict
from collections import OrderedDict

def solution(genres, plays):
    genre_dict = OrderedDict()  # 키: 장르 값: 리스트 (리스트 안에는 (재생 횟수, 노래id))
    count_dict = defaultdict(int)

    for i in range(len(genres)):
        if genres[i] in genre_dict.keys():
            genre_dict[genres[i]].append((plays[i], i))
        else:
            genre_dict[genres[i]] = list()
            genre_dict[genres[i]].append((plays[i], i))

        count_dict[genres[i]] += plays[i]

    count_dict = sorted(count_dict.items(), key = lambda x: x[1], reverse=True)

    answer = []
    for genre in count_dict:
        plays_id = genre_dict[genre[0]]
        plays_id.sort(key = lambda x: (-x[0], x[1]))

        if len(plays_id) >= 2:
            answer.append(plays_id[0][1])
            answer.append(plays_id[1][1])
        elif len(plays_id) == 1:
            answer.append(plays_id[0][1])

    return answer

두 개의 딕셔너리를 사용했다. 하나는 장르 : 장르 총 재생 횟수, 다른 하나는 장르 : ( 재생 횟수, 아이디).
딕셔너리를 정렬하는 방법을 기억해두자. sorted(dict.items(), )를 하면 튜플 형태로 반환된다.