알고리즘(Python)

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

문제링크

def solution(genres, plays):
    unique_genres = set(genres)
    dic = {}
    for gen in unique_genres:
        dic[gen] = []
    # 각 장르당 속하는 곡의 고유번호 리스트 만들기
    for genre in unique_genres:
        for i in range(len(genres)):
            if genre == genres[i]:
                dic[genre].append(i)

    def f1(x):
        lst = x[1]  # 해당 장르의 인덱스 리스트
        s = 0
        for i in lst:
            s += plays[i]
        return s
    # 장르마다 sum을 한 것을 기준으로 내림차순 정렬
    sort_dic = sorted(dic.items(), key=f1, reverse=True)
	# play수 내림차순, 고유번호 오름차순
    for i in sort_dic:
        i[1].sort(key= lambda x: (-plays[x], x))

    answer = []
    for i in sort_dic:
        for ii in i[1][:2]:
            answer.append(ii)
    return answer

코드가 겁나 긴데 통과하긴 했다.
sort함수에서 key를 잘 이용하는 것이 관건이었다.