[Python] 프로그래머스 / 베스트앨범

문제 링크

 

뭔가 좀 기교가 가득해보이는 코드... 흑백요리사도 아니고 코딩테스트인데 그게 나쁜 건 아니지만.

from collections import defaultdict

def solution(genres, plays):
    
    genre_cnt = defaultdict(list)
    genre_sum = defaultdict(int)
    
    for i, zip_ in enumerate(zip(genres, plays)) :
        genre, play = zip_
        genre_cnt[genre].append((play, i))
        genre_sum[genre] += play
        
    sorted_genres = sorted(genre_sum, key=lambda x: -genre_sum[x])
    
    for _, value in genre_cnt.items() :
        value.sort(key = lambda x: (-x[0], x[1]))

    answer = []
    for sorted_genre in sorted_genres :
        _, idx = genre_cnt[sorted_genre][0]
        answer.append(idx)
        
        if len(genre_cnt[sorted_genre]) > 1 :
            _, idx = genre_cnt[sorted_genre][1]
            answer.append(idx)
            
    return answer