[Algorithm] Programmers_42579_베스트앨범 - Go

문제 링크

Programmers_42579_베스트앨범

생각

1. 장르의 우선순위를 먼저 생각

2. 장르별로 노래의 우선순위를 생각

기본적으로 sort만을 사용해서도 풀이가 가능하지만 우선순위 큐를 체득하고 싶어서 우선순위 큐도 사용

코드

import (
   "container/heap"
   "fmt"
   "sort"
)

type genre struct {
   name       string
   playAmount int
}

type song struct {
   index     int
   playCount int
}

type pq []song

func (p pq) Len() int { return len(p) }

func (p pq) Less(i, j int) bool {
   if p[i].playCount == p[j].playCount {
      return p[i].index < p[j].index
   }
   return p[i].playCount > p[j].playCount
}

func (p pq) Swap(i, j int) { p[i], p[j] = p[j], p[i] }

func (p *pq) Push(x interface{}) {
   *p = append(*p, x.(song))
}

func (p *pq) Pop() interface{} {
   len := len(*p) - 1
   val := (*p)[len]
   *p = (*p)[0:len]
   return val
}

type genreChart []genre

func (g genreChart) Less(i, j int) bool { return g[i].playAmount > g[j].playAmount }
func (g genreChart) Swap(i, j int)      { g[i], g[j] = g[j], g[i] }
func (g genreChart) Len() int           { return len(g) }

func initGenreChart(genres []string, plays []int) genreChart {
   genreTable := make(map[string]int)
   genreChart := make(genreChart, 0)
   for i := range genres {
      if _, ok := genreTable[genres[i]]; !ok {
         genreTable[genres[i]] = plays[i]
         continue
      }
      genreTable[genres[i]] += plays[i]
   }
   for g := range genreTable {
      genreChart = append(genreChart, genre{g, genreTable[g]})
   }
   sort.Sort(genreChart)

   return genreChart
}

func initSongTable(genres []string, plays []int) map[string]*pq {
   songTable := make(map[string]*pq)
   for i, _ := range genres {
      if _, ok := songTable[genres[i]]; !ok {
         p := make(pq, 0)
         songTable[genres[i]] = &p
      }
      *songTable[genres[i]] = append(*songTable[genres[i]], song{i, plays[i]})
   }
   for key := range songTable {
      heap.Init(songTable[key])
   }
   return songTable
}

func makeSongRanking(gc genreChart, st map[string]*pq) []int {
   var ans []int
   for _, g := range gc {
      i := 0
      for true {
         ans = append(ans, heap.Pop(st[g.name]).(song).index)
         i++
         if st[g.name].Len() == 0 || i == 2 {
            break
         }
      }
   }
   return ans
}

func Solution(genres []string, plays []int) []int {
   genreChart := initGenreChart(genres, plays)
   songTable := initSongTable(genres, plays)

   fmt.Println(genreChart)
   fmt.Println(songTable)
   return makeSongRanking(genreChart, songTable)
}

 

결과