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

코딩테스트 연습 - 베스트앨범

import sys
from collections import defaultdict

input = sys.stdin.readline

genres = list(map(str, input().split()))
plays = list(map(int, input().split()))


def solution(genres, plays):
    answer = []
    genresPlays = defaultdict(int)
    genresSongs = defaultdict(lambda: [])
    i = 0
    for g, p in zip(genres, plays):
        genresPlays[g] += p
        genresSongs[g].append((i, p))
        i += 1

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

    for g in sortedGenres:
        sortedG = sorted(genresSongs[g[0]], key=lambda x: x[1], reverse=True)
        answer.append(sortedG[0][0])
        if len(sortedG) > 1:
            answer.append(sortedG[1][0])

    return answer


print(solution(genres, plays))

 

해설

해설에 앞서 defaultdict에 대해 설명하겠다. defaultdict란 dict에 없는 값을 넣을 경우 생기는 에러를 위해 쓰는 라이브러리다. 기존 dict에 없는 값을 넣으려면 조건문을 사용해야 했는데 defaultdict를 이용해 조건문 없이 더 간결한 코드로 기존 dict에 없는 값을 넣을 수 있다.

 

genresPlays에는 각 장르 별로 재생횟수 합을 저장한다.

genresSongs에는 각 장르에 들어간 곡의 인덱스와 재생횟수를 저장한다.

 

genresPlays를 재생횟수가 많은 장르 순으로 정렬한다.

 

genresPlays는 내림차순으로 정렬이 되어있다. 이를 이용해 genresSongs에서 그 장르 순서대로 장르의 각 곡들의 재생횟수를 기준으로 내림차순 정렬한다. 하나의 장르가 정렬이 완료되면 곡들의 재생횟수가 많은 곡부터 인덱스를 answer에 저장한다. 최대 2개 씩 노래를 저장할 수 있기 때문에 storedG의 길이가 2 이상이면 두 개를 모두, 아니면 한 개만 answer에 저장한다.