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

https://school.programmers.co.kr/learn/courses/30/lessons/42579?language=python3

프로그래머스

풀이

functools 모듈의 cmp_to_key 함수를 사용해 비교함수를 작성하여 답을 구할 수 있습니다.

 

sorted 함수를 사용해 정렬하는 경우 리턴은 list 타입임에 유의합니다. 

 

장르 우선 순위를 결정하기 위한 딕셔너리 dic과 모든 곡을 기록하고 비교하기 위한 album_list를 선언합니다. 

 

dic의 경우 값을 기준으로 내림차순 정렬하고, album_list는 cmp_to_key를 사용해 커스텀 정렬함수를 작성하여 문제에서 정의한대로 정렬해줍니다.

 

이후 dic의 가장 처음 원소의 장르부터 album_list의 앞 원소부터 선택하여 인덱스를 기록해줍니다. 이 때, 2곡을 수록했다면 break로 빠져나가줍니다. 

from functools import cmp_to_key

def album_compare(album1, album2):
    if album1.play < album2.play:
        return 1
    elif album1.play == album2.play:
        if album1.track < album2.track:
            return -1
        else:
            return 1
    else:
        return -1
    
def solution(genres, plays):
    answer = []
    dic = {}
    album_list = []
    for idx, genre in enumerate(genres):
        dic[genre] = dic.get(genre, 0) + plays[idx]
        album_list.append(album(genre, plays[idx], idx))

    dic = sorted(dic.items(), key=lambda dic:dic[1], reverse=True)
    album_list = sorted(album_list, key=cmp_to_key(album_compare))
    # print(dic)

    while len(dic) > 0:
        # 정렬했으므로 dic은 list 타입이다
        play_genre = dic.pop(0)
        # print(play_genre)
        cnt = 0;
        for ab in album_list:
            if play_genre[0] == ab.genre:
                answer.append(ab.track)
                cnt += 1
            if cnt == 2:
                break

    return answer

class album:
    def __init__(self, genre, play, track):
        self.genre = genre
        self.play = play
        self.track = track